From a0c57f4ca33d017600176feb6d5732468378fc03 Mon Sep 17 00:00:00 2001 From: viclafargue Date: Fri, 3 Feb 2023 16:31:58 +0100 Subject: [PATCH 01/20] Allow use of mdspan view in IVF-PQ API --- cpp/bench/neighbors/knn.cuh | 12 +++- cpp/include/raft/neighbors/ivf_pq.cuh | 48 ++++++------- .../raft/spatial/knn/detail/ann_quantized.cuh | 13 +++- cpp/include/raft_runtime/neighbors/ivf_pq.hpp | 65 ++++++++---------- cpp/src/distance/neighbors/ivfpq_build.cu | 67 +++++++++---------- cpp/src/distance/neighbors/ivfpq_search.cu | 25 ++++--- cpp/test/neighbors/ann_ivf_pq.cuh | 35 ++++++---- 7 files changed, 137 insertions(+), 128 deletions(-) diff --git a/cpp/bench/neighbors/knn.cuh b/cpp/bench/neighbors/knn.cuh index eec1cba99e..7a2eadd096 100644 --- a/cpp/bench/neighbors/knn.cuh +++ b/cpp/bench/neighbors/knn.cuh @@ -183,8 +183,9 @@ struct ivf_pq_knn { { index_params.n_lists = 4096; index_params.metric = raft::distance::DistanceType::L2Expanded; - index.emplace(raft::neighbors::ivf_pq::build( - handle, index_params, data, IdxT(ps.n_samples), uint32_t(ps.n_dims))); + + auto data_view = raft::make_device_matrix_view(data, ps.n_samples, ps.n_dims); + index.emplace(raft::neighbors::ivf_pq::build(handle, index_params, data_view)); } void search(const raft::device_resources& handle, @@ -193,8 +194,13 @@ struct ivf_pq_knn { IdxT* out_idxs) { search_params.n_probes = 20; + + auto queries_view = + raft::make_device_matrix_view(search_items, ps.n_queries, ps.n_dims); + auto idxs_view = raft::make_device_matrix_view(out_idxs, ps.n_queries, ps.k); + auto dists_view = raft::make_device_matrix_view(out_dists, ps.n_queries, ps.k); raft::neighbors::ivf_pq::search( - handle, search_params, *index, search_items, ps.n_queries, ps.k, out_idxs, out_dists); + handle, search_params, *index, queries_view, ps.k, idxs_view, dists_view); } }; diff --git a/cpp/include/raft/neighbors/ivf_pq.cuh b/cpp/include/raft/neighbors/ivf_pq.cuh index 287f0bc5f4..fe1817c203 100644 --- a/cpp/include/raft/neighbors/ivf_pq.cuh +++ b/cpp/include/raft/neighbors/ivf_pq.cuh @@ -59,19 +59,18 @@ namespace raft::neighbors::ivf_pq { * @param handle * @param params configure the index building * @param[in] dataset a device/host pointer to a row-major matrix [n_rows, dim] - * @param n_rows the number of samples - * @param dim the dimensionality of the data * * @return the constructed ivf-pq index */ template inline auto build(raft::device_resources const& handle, const index_params& params, - const T* dataset, - IdxT n_rows, - uint32_t dim) -> index + const raft::device_matrix_view& dataset) -> index { - return raft::spatial::knn::ivf_pq::detail::build(handle, params, dataset, n_rows, dim); + IdxT n_rows = dataset.extent(0); + IdxT dim = dataset.extent(1); + return raft::spatial::knn::ivf_pq::detail::build( + handle, params, dataset.data_handle(), n_rows, dim); } /** @@ -102,19 +101,18 @@ inline auto build(raft::device_resources const& handle, * @param[in] new_indices a device/host pointer to a vector of indices [n_rows]. * If the original index is empty (`orig_index.size() == 0`), you can pass `nullptr` * here to imply a continuous range `[0...n_rows)`. - * @param n_rows the number of samples * * @return the constructed extended ivf-pq index */ template inline auto extend(raft::device_resources const& handle, const index& orig_index, - const T* new_vectors, - const IdxT* new_indices, - IdxT n_rows) -> index + const raft::device_matrix_view& new_vectors, + const raft::device_matrix_view& new_indices) -> index { + IdxT n_rows = new_vectors.extent(0); return raft::spatial::knn::ivf_pq::detail::extend( - handle, orig_index, new_vectors, new_indices, n_rows); + handle, orig_index, new_vectors.data_handle(), new_indices.data_handle(), n_rows); } /** @@ -129,16 +127,14 @@ inline auto extend(raft::device_resources const& handle, * @param[in] new_indices a device/host pointer to a vector of indices [n_rows]. * If the original index is empty (`orig_index.size() == 0`), you can pass `nullptr` * here to imply a continuous range `[0...n_rows)`. - * @param n_rows the number of samples */ template inline void extend(raft::device_resources const& handle, index* index, - const T* new_vectors, - const IdxT* new_indices, - IdxT n_rows) + const raft::device_matrix_view& new_vectors, + const raft::device_matrix_view& new_indices) { - *index = extend(handle, *index, new_vectors, new_indices, n_rows); + *index = extend(handle, *index, new_vectors, new_indices); } /** @@ -175,7 +171,6 @@ inline void extend(raft::device_resources const& handle, * @param params configure the search * @param index ivf-pq constructed index * @param[in] queries a device pointer to a row-major matrix [n_queries, index->dim()] - * @param n_queries the batch size * @param k the number of neighbors to find for each query. * @param[out] neighbors a device pointer to the indices of the neighbors in the source dataset * [n_queries, k] @@ -187,15 +182,22 @@ template inline void search(raft::device_resources const& handle, const search_params& params, const index& index, - const T* queries, - uint32_t n_queries, + const raft::device_matrix_view& queries, uint32_t k, - IdxT* neighbors, - float* distances, + const raft::device_matrix_view& neighbors, + const raft::device_matrix_view& distances, rmm::mr::device_memory_resource* mr = nullptr) { - return raft::spatial::knn::ivf_pq::detail::search( - handle, params, index, queries, n_queries, k, neighbors, distances, mr); + IdxT n_queries = queries.extent(0); + return raft::spatial::knn::ivf_pq::detail::search(handle, + params, + index, + queries.data_handle(), + n_queries, + k, + neighbors.data_handle(), + distances.data_handle(), + mr); } /** @} */ // end group ivf_pq diff --git a/cpp/include/raft/spatial/knn/detail/ann_quantized.cuh b/cpp/include/raft/spatial/knn/detail/ann_quantized.cuh index 427e812cda..066dcaaa6b 100644 --- a/cpp/include/raft/spatial/knn/detail/ann_quantized.cuh +++ b/cpp/include/raft/spatial/knn/detail/ann_quantized.cuh @@ -78,8 +78,10 @@ void approx_knn_build_index(raft::device_resources const& handle, params.pq_bits = ivf_pq_pams->n_bits; params.pq_dim = ivf_pq_pams->M; // TODO: handle ivf_pq_pams.usePrecomputedTables ? - index->ivf_pq = std::make_unique>( - neighbors::ivf_pq::build(handle, params, index_array, int64_t(n), D)); + + auto index_view = raft::make_device_matrix_view(index_array, n, D); + index->ivf_pq = std::make_unique>( + neighbors::ivf_pq::build(handle, params, index_view)); } else { RAFT_FAIL("Unrecognized index type."); } @@ -110,8 +112,13 @@ void approx_knn_search(raft::device_resources const& handle, } else if (index->ivf_pq) { neighbors::ivf_pq::search_params params; params.n_probes = index->nprobe; + + auto query_view = + raft::make_device_matrix_view(query_array, n, index->ivf_pq->dim()); + auto indices_view = raft::make_device_matrix_view(indices, n, k); + auto distances_view = raft::make_device_matrix_view(distances, n, k); neighbors::ivf_pq::search( - handle, params, *index->ivf_pq, query_array, n, k, indices, distances); + handle, params, *index->ivf_pq, query_view, k, indices_view, distances_view); } else { RAFT_FAIL("The model is not trained"); } diff --git a/cpp/include/raft_runtime/neighbors/ivf_pq.hpp b/cpp/include/raft_runtime/neighbors/ivf_pq.hpp index cae32c9530..40272987c0 100644 --- a/cpp/include/raft_runtime/neighbors/ivf_pq.hpp +++ b/cpp/include/raft_runtime/neighbors/ivf_pq.hpp @@ -20,15 +20,14 @@ namespace raft::runtime::neighbors::ivf_pq { -#define RAFT_INST_SEARCH(T, IdxT) \ - void search(raft::device_resources const&, \ - const raft::neighbors::ivf_pq::search_params&, \ - const raft::neighbors::ivf_pq::index&, \ - const T*, \ - uint32_t, \ - uint32_t, \ - IdxT*, \ - float*, \ +#define RAFT_INST_SEARCH(T, IdxT) \ + void search(raft::device_resources const&, \ + const raft::neighbors::ivf_pq::search_params&, \ + const raft::neighbors::ivf_pq::index&, \ + const raft::device_matrix_view&, \ + uint32_t, \ + const raft::device_matrix_view&, \ + const raft::device_matrix_view&, \ rmm::mr::device_memory_resource*); RAFT_INST_SEARCH(float, uint64_t); @@ -40,33 +39,27 @@ RAFT_INST_SEARCH(uint8_t, uint64_t); // We define overloads for build and extend with void return type. This is used in the Cython // wrappers, where exception handling is not compatible with return type that has nontrivial // constructor. -#define RAFT_INST_BUILD_EXTEND(T, IdxT) \ - auto build(raft::device_resources const& handle, \ - const raft::neighbors::ivf_pq::index_params& params, \ - const T* dataset, \ - IdxT n_rows, \ - uint32_t dim) \ - ->raft::neighbors::ivf_pq::index; \ - \ - auto extend(raft::device_resources const& handle, \ - const raft::neighbors::ivf_pq::index& orig_index, \ - const T* new_vectors, \ - const IdxT* new_indices, \ - IdxT n_rows) \ - ->raft::neighbors::ivf_pq::index; \ - \ - void build(raft::device_resources const& handle, \ - const raft::neighbors::ivf_pq::index_params& params, \ - const T* dataset, \ - IdxT n_rows, \ - uint32_t dim, \ - raft::neighbors::ivf_pq::index* idx); \ - \ - void extend(raft::device_resources const& handle, \ - raft::neighbors::ivf_pq::index* idx, \ - const T* new_vectors, \ - const IdxT* new_indices, \ - IdxT n_rows); +#define RAFT_INST_BUILD_EXTEND(T, IdxT) \ + auto build(raft::device_resources const& handle, \ + const raft::neighbors::ivf_pq::index_params& params, \ + const raft::device_matrix_view& dataset) \ + ->raft::neighbors::ivf_pq::index; \ + \ + auto extend(raft::device_resources const& handle, \ + const raft::neighbors::ivf_pq::index& orig_index, \ + const raft::device_matrix_view& new_vectors, \ + const raft::device_matrix_view& new_indices) \ + ->raft::neighbors::ivf_pq::index; \ + \ + void build(raft::device_resources const& handle, \ + const raft::neighbors::ivf_pq::index_params& params, \ + const raft::device_matrix_view& dataset, \ + raft::neighbors::ivf_pq::index* idx); \ + \ + void extend(raft::device_resources const& handle, \ + raft::neighbors::ivf_pq::index* idx, \ + const raft::device_matrix_view& new_vectors, \ + const raft::device_matrix_view& new_indices); RAFT_INST_BUILD_EXTEND(float, uint64_t) RAFT_INST_BUILD_EXTEND(int8_t, uint64_t) diff --git a/cpp/src/distance/neighbors/ivfpq_build.cu b/cpp/src/distance/neighbors/ivfpq_build.cu index 650767f918..caa2092543 100644 --- a/cpp/src/distance/neighbors/ivfpq_build.cu +++ b/cpp/src/distance/neighbors/ivfpq_build.cu @@ -19,43 +19,36 @@ namespace raft::runtime::neighbors::ivf_pq { -#define RAFT_INST_BUILD_EXTEND(T, IdxT) \ - auto build(raft::device_resources const& handle, \ - const raft::neighbors::ivf_pq::index_params& params, \ - const T* dataset, \ - IdxT n_rows, \ - uint32_t dim) \ - ->raft::neighbors::ivf_pq::index \ - { \ - return raft::neighbors::ivf_pq::build(handle, params, dataset, n_rows, dim); \ - } \ - auto extend(raft::device_resources const& handle, \ - const raft::neighbors::ivf_pq::index& orig_index, \ - const T* new_vectors, \ - const IdxT* new_indices, \ - IdxT n_rows) \ - ->raft::neighbors::ivf_pq::index \ - { \ - return raft::neighbors::ivf_pq::extend( \ - handle, orig_index, new_vectors, new_indices, n_rows); \ - } \ - \ - void build(raft::device_resources const& handle, \ - const raft::neighbors::ivf_pq::index_params& params, \ - const T* dataset, \ - IdxT n_rows, \ - uint32_t dim, \ - raft::neighbors::ivf_pq::index* idx) \ - { \ - *idx = raft::neighbors::ivf_pq::build(handle, params, dataset, n_rows, dim); \ - } \ - void extend(raft::device_resources const& handle, \ - raft::neighbors::ivf_pq::index* idx, \ - const T* new_vectors, \ - const IdxT* new_indices, \ - IdxT n_rows) \ - { \ - raft::neighbors::ivf_pq::extend(handle, idx, new_vectors, new_indices, n_rows); \ +#define RAFT_INST_BUILD_EXTEND(T, IdxT) \ + auto build(raft::device_resources const& handle, \ + const raft::neighbors::ivf_pq::index_params& params, \ + const raft::device_matrix_view& dataset) \ + ->raft::neighbors::ivf_pq::index \ + { \ + return raft::neighbors::ivf_pq::build(handle, params, dataset); \ + } \ + auto extend(raft::device_resources const& handle, \ + const raft::neighbors::ivf_pq::index& orig_index, \ + const raft::device_matrix_view& new_vectors, \ + const raft::device_matrix_view& new_indices) \ + ->raft::neighbors::ivf_pq::index \ + { \ + return raft::neighbors::ivf_pq::extend(handle, orig_index, new_vectors, new_indices); \ + } \ + \ + void build(raft::device_resources const& handle, \ + const raft::neighbors::ivf_pq::index_params& params, \ + const raft::device_matrix_view& dataset, \ + raft::neighbors::ivf_pq::index* idx) \ + { \ + *idx = raft::neighbors::ivf_pq::build(handle, params, dataset); \ + } \ + void extend(raft::device_resources const& handle, \ + raft::neighbors::ivf_pq::index* idx, \ + const raft::device_matrix_view& new_vectors, \ + const raft::device_matrix_view& new_indices) \ + { \ + raft::neighbors::ivf_pq::extend(handle, idx, new_vectors, new_indices); \ } RAFT_INST_BUILD_EXTEND(float, uint64_t); diff --git a/cpp/src/distance/neighbors/ivfpq_search.cu b/cpp/src/distance/neighbors/ivfpq_search.cu index 05ab890ea5..6ce4cd8375 100644 --- a/cpp/src/distance/neighbors/ivfpq_search.cu +++ b/cpp/src/distance/neighbors/ivfpq_search.cu @@ -20,19 +20,18 @@ namespace raft::runtime::neighbors::ivf_pq { -#define RAFT_SEARCH_INST(T, IdxT) \ - void search(raft::device_resources const& handle, \ - const raft::neighbors::ivf_pq::search_params& params, \ - const raft::neighbors::ivf_pq::index& idx, \ - const T* queries, \ - uint32_t n_queries, \ - uint32_t k, \ - IdxT* neighbors, \ - float* distances, \ - rmm::mr::device_memory_resource* mr) \ - { \ - raft::neighbors::ivf_pq::search( \ - handle, params, idx, queries, n_queries, k, neighbors, distances, mr); \ +#define RAFT_SEARCH_INST(T, IdxT) \ + void search(raft::device_resources const& handle, \ + const raft::neighbors::ivf_pq::search_params& params, \ + const raft::neighbors::ivf_pq::index& idx, \ + const raft::device_matrix_view& queries, \ + uint32_t k, \ + const raft::device_matrix_view& neighbors, \ + const raft::device_matrix_view& distances, \ + rmm::mr::device_memory_resource* mr) \ + { \ + raft::neighbors::ivf_pq::search( \ + handle, params, idx, queries, k, neighbors, distances, mr); \ } RAFT_SEARCH_INST(float, uint64_t); diff --git a/cpp/test/neighbors/ann_ivf_pq.cuh b/cpp/test/neighbors/ann_ivf_pq.cuh index 488041f527..bf60c2f589 100644 --- a/cpp/test/neighbors/ann_ivf_pq.cuh +++ b/cpp/test/neighbors/ann_ivf_pq.cuh @@ -183,7 +183,9 @@ class ivf_pq_test : public ::testing::TestWithParam { auto ipams = ps.index_params; ipams.add_data_on_build = true; - return ivf_pq::build(handle_, ipams, database.data(), ps.num_db_vecs, ps.dim); + auto index_view = + raft::make_device_matrix_view(database.data(), ps.num_db_vecs, ps.dim); + return ivf_pq::build(handle_, ipams, index_view); } auto build_2_extends() @@ -203,11 +205,17 @@ class ivf_pq_test : public ::testing::TestWithParam { auto ipams = ps.index_params; ipams.add_data_on_build = false; - auto index = - ivf_pq::build(handle_, ipams, database.data(), ps.num_db_vecs, ps.dim); + auto database_view = + raft::make_device_matrix_view(database.data(), ps.num_db_vecs, ps.dim); + auto index = ivf_pq::build(handle_, ipams, database_view); - ivf_pq::extend(handle_, &index, vecs_2, inds_2, size_2); - return ivf_pq::extend(handle_, index, vecs_1, inds_1, size_1); + auto vecs_2_view = raft::make_device_matrix_view(vecs_2, size_2, ps.dim); + auto inds_2_view = raft::make_device_matrix_view(inds_2, size_2, 1); + ivf_pq::extend(handle_, &index, vecs_2_view, inds_2_view); + + auto vecs_1_view = raft::make_device_matrix_view(vecs_1, size_1, ps.dim); + auto inds_1_view = raft::make_device_matrix_view(inds_1, size_1, 1); + return ivf_pq::extend(handle_, index, vecs_1_view, inds_1_view); } template @@ -226,14 +234,15 @@ class ivf_pq_test : public ::testing::TestWithParam { rmm::device_uvector distances_ivf_pq_dev(queries_size, stream_); rmm::device_uvector indices_ivf_pq_dev(queries_size, stream_); - ivf_pq::search(handle_, - ps.search_params, - index, - search_queries.data(), - ps.num_queries, - ps.k, - indices_ivf_pq_dev.data(), - distances_ivf_pq_dev.data()); + auto query_view = + raft::make_device_matrix_view(search_queries.data(), ps.num_queries, ps.dim); + auto inds_view = + raft::make_device_matrix_view(indices_ivf_pq_dev.data(), ps.num_queries, ps.k); + auto dists_view = + raft::make_device_matrix_view(distances_ivf_pq_dev.data(), ps.num_queries, ps.k); + + ivf_pq::search( + handle_, ps.search_params, index, query_view, ps.k, inds_view, dists_view); update_host(distances_ivf_pq.data(), distances_ivf_pq_dev.data(), queries_size, stream_); update_host(indices_ivf_pq.data(), indices_ivf_pq_dev.data(), queries_size, stream_); From 45553a1b3364744ae7d7ceb1a639bc8f2a93fef7 Mon Sep 17 00:00:00 2001 From: viclafargue Date: Mon, 6 Feb 2023 11:12:29 +0100 Subject: [PATCH 02/20] restore legacy API --- cpp/include/raft/neighbors/ivf_pq.cuh | 128 ++++++++++++++++++ cpp/include/raft_runtime/neighbors/ivf_pq.hpp | 36 +++++ 2 files changed, 164 insertions(+) diff --git a/cpp/include/raft/neighbors/ivf_pq.cuh b/cpp/include/raft/neighbors/ivf_pq.cuh index fe1817c203..27f0404b6d 100644 --- a/cpp/include/raft/neighbors/ivf_pq.cuh +++ b/cpp/include/raft/neighbors/ivf_pq.cuh @@ -59,6 +59,35 @@ namespace raft::neighbors::ivf_pq { * @param handle * @param params configure the index building * @param[in] dataset a device/host pointer to a row-major matrix [n_rows, dim] + * @param n_rows the number of samples + * @param dim the dimensionality of the data + * + * @return the constructed ivf-pq index + */ +template +inline auto build(raft::device_resources const& handle, + const index_params& params, + const T* dataset, + IdxT n_rows, + uint32_t dim) -> index +{ + return raft::spatial::knn::ivf_pq::detail::build(handle, params, dataset, n_rows, dim); +} + +/** + * @brief Build the index from the dataset for efficient search. + * + * NB: Currently, the following distance metrics are supported: + * - L2Expanded + * - L2Unexpanded + * - InnerProduct + * + * @tparam T data element type + * @tparam IdxT type of the indices in the source dataset + * + * @param handle + * @param 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 */ @@ -101,6 +130,37 @@ inline auto build(raft::device_resources const& handle, * @param[in] new_indices a device/host pointer to a vector of indices [n_rows]. * If the original index is empty (`orig_index.size() == 0`), you can pass `nullptr` * here to imply a continuous range `[0...n_rows)`. + * @param n_rows the number of samples + * + * @return the constructed extended ivf-pq index + */ +template +inline auto extend(raft::device_resources const& handle, + const index& orig_index, + const T* new_vectors, + const IdxT* new_indices, + IdxT n_rows) -> index +{ + return raft::spatial::knn::ivf_pq::detail::extend( + handle, orig_index, new_vectors, new_indices, n_rows); +} + +/** + * @brief Build a new index containing the data of the original plus new extra vectors. + * + * Implementation note: + * The new data is clustered according to existing kmeans clusters, then the cluster + * centers are unchanged. + * + * @tparam T data element type + * @tparam IdxT type of the indices in the source dataset + * + * @param handle + * @param orig_index original index + * @param[in] new_vectors a device matrix view to a row-major matrix [n_rows, index.dim()] + * @param[in] new_indices a device matrix view to a vector of indices [n_rows]. + * If the original index is empty (`orig_index.size() == 0`), you can pass `nullptr` + * here to imply a continuous range `[0...n_rows)`. * * @return the constructed extended ivf-pq index */ @@ -127,6 +187,30 @@ inline auto extend(raft::device_resources const& handle, * @param[in] new_indices a device/host pointer to a vector of indices [n_rows]. * If the original index is empty (`orig_index.size() == 0`), you can pass `nullptr` * here to imply a continuous range `[0...n_rows)`. + * @param n_rows the number of samples + */ +template +inline void extend(raft::device_resources const& handle, + index* index, + const T* new_vectors, + const IdxT* new_indices, + IdxT n_rows) +{ + *index = extend(handle, *index, new_vectors, new_indices, n_rows); +} + +/** + * @brief Extend the index with the new data. + * * + * @tparam T data element type + * @tparam IdxT type of the indices in the source dataset + * + * @param handle + * @param[inout] index + * @param[in] new_vectors a device matrix view to a row-major matrix [n_rows, index.dim()] + * @param[in] new_indices a device matrix view to a vector of indices [n_rows]. + * If the original index is empty (`orig_index.size() == 0`), you can pass `nullptr` + * here to imply a continuous range `[0...n_rows)`. */ template inline void extend(raft::device_resources const& handle, @@ -171,6 +255,7 @@ inline void extend(raft::device_resources const& handle, * @param params configure the search * @param index ivf-pq constructed index * @param[in] queries a device pointer to a row-major matrix [n_queries, index->dim()] + * @param n_queries the batch size * @param k the number of neighbors to find for each query. * @param[out] neighbors a device pointer to the indices of the neighbors in the source dataset * [n_queries, k] @@ -179,6 +264,49 @@ inline void extend(raft::device_resources const& handle, * memory pool here to avoid memory allocations within search). */ template +inline void search(raft::device_resources const& handle, + const search_params& params, + const index& index, + const T* queries, + uint32_t n_queries, + uint32_t k, + IdxT* neighbors, + float* distances, + rmm::mr::device_memory_resource* mr = nullptr) +{ + return raft::spatial::knn::ivf_pq::detail::search( + handle, params, index, queries, n_queries, k, neighbors, distances, mr); +} + +/** + * @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`. + * + * @tparam T data element type + * @tparam IdxT type of the indices + * + * @param handle + * @param params configure the search + * @param index ivf-pq constructed index + * @param[in] queries a device matrix view to a row-major matrix [n_queries, index->dim()] + * @param k the number of neighbors to find for each query. + * @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 mr an optional memory resource to use across the searches (you can provide a large enough + * memory pool here to avoid memory allocations within search). + */ +template inline void search(raft::device_resources const& handle, const search_params& params, const index& index, diff --git a/cpp/include/raft_runtime/neighbors/ivf_pq.hpp b/cpp/include/raft_runtime/neighbors/ivf_pq.hpp index 40272987c0..7e83aa668a 100644 --- a/cpp/include/raft_runtime/neighbors/ivf_pq.hpp +++ b/cpp/include/raft_runtime/neighbors/ivf_pq.hpp @@ -21,6 +21,16 @@ namespace raft::runtime::neighbors::ivf_pq { #define RAFT_INST_SEARCH(T, IdxT) \ + void search(raft::device_resources const&, \ + const raft::neighbors::ivf_pq::search_params&, \ + const raft::neighbors::ivf_pq::index&, \ + const T*, \ + uint32_t, \ + uint32_t, \ + IdxT*, \ + float*, \ + rmm::mr::device_memory_resource*); \ + \ void search(raft::device_resources const&, \ const raft::neighbors::ivf_pq::search_params&, \ const raft::neighbors::ivf_pq::index&, \ @@ -40,6 +50,32 @@ RAFT_INST_SEARCH(uint8_t, uint64_t); // wrappers, where exception handling is not compatible with return type that has nontrivial // constructor. #define RAFT_INST_BUILD_EXTEND(T, IdxT) \ + auto build(raft::device_resources const& handle, \ + const raft::neighbors::ivf_pq::index_params& params, \ + const T* dataset, \ + IdxT n_rows, \ + uint32_t dim) \ + ->raft::neighbors::ivf_pq::index; \ + \ + auto extend(raft::device_resources const& handle, \ + const raft::neighbors::ivf_pq::index& orig_index, \ + const T* new_vectors, \ + const IdxT* new_indices, \ + IdxT n_rows) \ + ->raft::neighbors::ivf_pq::index; \ + \ + void build(raft::device_resources const& handle, \ + const raft::neighbors::ivf_pq::index_params& params, \ + const T* dataset, \ + IdxT n_rows, \ + uint32_t dim, \ + raft::neighbors::ivf_pq::index* idx); \ + \ + void extend(raft::device_resources const& handle, \ + raft::neighbors::ivf_pq::index* idx, \ + const T* new_vectors, \ + const IdxT* new_indices, \ + IdxT n_rows); \ auto build(raft::device_resources const& handle, \ const raft::neighbors::ivf_pq::index_params& params, \ const raft::device_matrix_view& dataset) \ From c602e8d02f49b2a5055824c477dee271e5f66028 Mon Sep 17 00:00:00 2001 From: viclafargue Date: Mon, 6 Feb 2023 14:47:00 +0100 Subject: [PATCH 03/20] row_major + assert tests --- cpp/include/raft/neighbors/ivf_pq.cuh | 24 ++-- cpp/include/raft_runtime/neighbors/ivf_pq.hpp | 130 +++++++++--------- cpp/src/distance/neighbors/ivfpq_build.cu | 12 +- cpp/src/distance/neighbors/ivfpq_search.cu | 24 ++-- 4 files changed, 98 insertions(+), 92 deletions(-) diff --git a/cpp/include/raft/neighbors/ivf_pq.cuh b/cpp/include/raft/neighbors/ivf_pq.cuh index 27f0404b6d..a945472e03 100644 --- a/cpp/include/raft/neighbors/ivf_pq.cuh +++ b/cpp/include/raft/neighbors/ivf_pq.cuh @@ -94,7 +94,7 @@ inline auto build(raft::device_resources const& handle, template inline auto build(raft::device_resources const& handle, const index_params& params, - const raft::device_matrix_view& dataset) -> index + const raft::device_matrix_view& dataset) -> index { IdxT n_rows = dataset.extent(0); IdxT dim = dataset.extent(1); @@ -167,10 +167,13 @@ inline auto extend(raft::device_resources const& handle, template inline auto extend(raft::device_resources const& handle, const index& orig_index, - const raft::device_matrix_view& new_vectors, - const raft::device_matrix_view& new_indices) -> index + const raft::device_matrix_view& new_vectors, + const raft::device_matrix_view& new_indices) + -> index { IdxT n_rows = new_vectors.extent(0); + ASSERT(n_rows == new_indices.extent(0), + "new_vectors and new_indices have different number of rows"); return raft::spatial::knn::ivf_pq::detail::extend( handle, orig_index, new_vectors.data_handle(), new_indices.data_handle(), n_rows); } @@ -215,8 +218,8 @@ inline void extend(raft::device_resources const& handle, template inline void extend(raft::device_resources const& handle, index* index, - const raft::device_matrix_view& new_vectors, - const raft::device_matrix_view& new_indices) + const raft::device_matrix_view& new_vectors, + const raft::device_matrix_view& new_indices) { *index = extend(handle, *index, new_vectors, new_indices); } @@ -310,13 +313,16 @@ template inline void search(raft::device_resources const& handle, const search_params& params, const index& index, - const raft::device_matrix_view& queries, + const raft::device_matrix_view& queries, uint32_t k, - const raft::device_matrix_view& neighbors, - const raft::device_matrix_view& distances, + const raft::device_matrix_view& neighbors, + const raft::device_matrix_view& distances, rmm::mr::device_memory_resource* mr = nullptr) { - IdxT n_queries = queries.extent(0); + IdxT n_queries = queries.extent(0); + bool check_n_rows = (n_queries == neighbors.extent(0)) && (n_queries == distances.extent(0)); + ASSERT(check_n_rows, + "queries, neighbors and distances parameters have inconsistent number of rows"); return raft::spatial::knn::ivf_pq::detail::search(handle, params, index, diff --git a/cpp/include/raft_runtime/neighbors/ivf_pq.hpp b/cpp/include/raft_runtime/neighbors/ivf_pq.hpp index 7e83aa668a..8745d1e568 100644 --- a/cpp/include/raft_runtime/neighbors/ivf_pq.hpp +++ b/cpp/include/raft_runtime/neighbors/ivf_pq.hpp @@ -20,24 +20,24 @@ namespace raft::runtime::neighbors::ivf_pq { -#define RAFT_INST_SEARCH(T, IdxT) \ - void search(raft::device_resources const&, \ - const raft::neighbors::ivf_pq::search_params&, \ - const raft::neighbors::ivf_pq::index&, \ - const T*, \ - uint32_t, \ - uint32_t, \ - IdxT*, \ - float*, \ - rmm::mr::device_memory_resource*); \ - \ - void search(raft::device_resources const&, \ - const raft::neighbors::ivf_pq::search_params&, \ - const raft::neighbors::ivf_pq::index&, \ - const raft::device_matrix_view&, \ - uint32_t, \ - const raft::device_matrix_view&, \ - const raft::device_matrix_view&, \ +#define RAFT_INST_SEARCH(T, IdxT) \ + void search(raft::device_resources const&, \ + const raft::neighbors::ivf_pq::search_params&, \ + const raft::neighbors::ivf_pq::index&, \ + const T*, \ + uint32_t, \ + uint32_t, \ + IdxT*, \ + float*, \ + rmm::mr::device_memory_resource*); \ + \ + void search(raft::device_resources const&, \ + const raft::neighbors::ivf_pq::search_params&, \ + const raft::neighbors::ivf_pq::index&, \ + const raft::device_matrix_view&, \ + uint32_t, \ + const raft::device_matrix_view&, \ + const raft::device_matrix_view&, \ rmm::mr::device_memory_resource*); RAFT_INST_SEARCH(float, uint64_t); @@ -49,53 +49,53 @@ RAFT_INST_SEARCH(uint8_t, uint64_t); // We define overloads for build and extend with void return type. This is used in the Cython // wrappers, where exception handling is not compatible with return type that has nontrivial // constructor. -#define RAFT_INST_BUILD_EXTEND(T, IdxT) \ - auto build(raft::device_resources const& handle, \ - const raft::neighbors::ivf_pq::index_params& params, \ - const T* dataset, \ - IdxT n_rows, \ - uint32_t dim) \ - ->raft::neighbors::ivf_pq::index; \ - \ - auto extend(raft::device_resources const& handle, \ - const raft::neighbors::ivf_pq::index& orig_index, \ - const T* new_vectors, \ - const IdxT* new_indices, \ - IdxT n_rows) \ - ->raft::neighbors::ivf_pq::index; \ - \ - void build(raft::device_resources const& handle, \ - const raft::neighbors::ivf_pq::index_params& params, \ - const T* dataset, \ - IdxT n_rows, \ - uint32_t dim, \ - raft::neighbors::ivf_pq::index* idx); \ - \ - void extend(raft::device_resources const& handle, \ - raft::neighbors::ivf_pq::index* idx, \ - const T* new_vectors, \ - const IdxT* new_indices, \ - IdxT n_rows); \ - auto build(raft::device_resources const& handle, \ - const raft::neighbors::ivf_pq::index_params& params, \ - const raft::device_matrix_view& dataset) \ - ->raft::neighbors::ivf_pq::index; \ - \ - auto extend(raft::device_resources const& handle, \ - const raft::neighbors::ivf_pq::index& orig_index, \ - const raft::device_matrix_view& new_vectors, \ - const raft::device_matrix_view& new_indices) \ - ->raft::neighbors::ivf_pq::index; \ - \ - void build(raft::device_resources const& handle, \ - const raft::neighbors::ivf_pq::index_params& params, \ - const raft::device_matrix_view& dataset, \ - raft::neighbors::ivf_pq::index* idx); \ - \ - void extend(raft::device_resources const& handle, \ - raft::neighbors::ivf_pq::index* idx, \ - const raft::device_matrix_view& new_vectors, \ - const raft::device_matrix_view& new_indices); +#define RAFT_INST_BUILD_EXTEND(T, IdxT) \ + auto build(raft::device_resources const& handle, \ + const raft::neighbors::ivf_pq::index_params& params, \ + const T* dataset, \ + IdxT n_rows, \ + uint32_t dim) \ + ->raft::neighbors::ivf_pq::index; \ + \ + auto extend(raft::device_resources const& handle, \ + const raft::neighbors::ivf_pq::index& orig_index, \ + const T* new_vectors, \ + const IdxT* new_indices, \ + IdxT n_rows) \ + ->raft::neighbors::ivf_pq::index; \ + \ + void build(raft::device_resources const& handle, \ + const raft::neighbors::ivf_pq::index_params& params, \ + const T* dataset, \ + IdxT n_rows, \ + uint32_t dim, \ + raft::neighbors::ivf_pq::index* idx); \ + \ + void extend(raft::device_resources const& handle, \ + raft::neighbors::ivf_pq::index* idx, \ + const T* new_vectors, \ + const IdxT* new_indices, \ + IdxT n_rows); \ + auto build(raft::device_resources const& handle, \ + const raft::neighbors::ivf_pq::index_params& params, \ + const raft::device_matrix_view& dataset) \ + ->raft::neighbors::ivf_pq::index; \ + \ + auto extend(raft::device_resources const& handle, \ + const raft::neighbors::ivf_pq::index& orig_index, \ + const raft::device_matrix_view& new_vectors, \ + const raft::device_matrix_view& new_indices) \ + ->raft::neighbors::ivf_pq::index; \ + \ + void build(raft::device_resources const& handle, \ + const raft::neighbors::ivf_pq::index_params& params, \ + const raft::device_matrix_view& dataset, \ + raft::neighbors::ivf_pq::index* idx); \ + \ + void extend(raft::device_resources const& handle, \ + raft::neighbors::ivf_pq::index* idx, \ + const raft::device_matrix_view& new_vectors, \ + const raft::device_matrix_view& new_indices); RAFT_INST_BUILD_EXTEND(float, uint64_t) RAFT_INST_BUILD_EXTEND(int8_t, uint64_t) diff --git a/cpp/src/distance/neighbors/ivfpq_build.cu b/cpp/src/distance/neighbors/ivfpq_build.cu index caa2092543..828f317908 100644 --- a/cpp/src/distance/neighbors/ivfpq_build.cu +++ b/cpp/src/distance/neighbors/ivfpq_build.cu @@ -22,15 +22,15 @@ namespace raft::runtime::neighbors::ivf_pq { #define RAFT_INST_BUILD_EXTEND(T, IdxT) \ auto build(raft::device_resources const& handle, \ const raft::neighbors::ivf_pq::index_params& params, \ - const raft::device_matrix_view& dataset) \ + const raft::device_matrix_view& dataset) \ ->raft::neighbors::ivf_pq::index \ { \ return raft::neighbors::ivf_pq::build(handle, params, dataset); \ } \ auto extend(raft::device_resources const& handle, \ const raft::neighbors::ivf_pq::index& orig_index, \ - const raft::device_matrix_view& new_vectors, \ - const raft::device_matrix_view& new_indices) \ + const raft::device_matrix_view& new_vectors, \ + const raft::device_matrix_view& new_indices) \ ->raft::neighbors::ivf_pq::index \ { \ return raft::neighbors::ivf_pq::extend(handle, orig_index, new_vectors, new_indices); \ @@ -38,15 +38,15 @@ namespace raft::runtime::neighbors::ivf_pq { \ void build(raft::device_resources const& handle, \ const raft::neighbors::ivf_pq::index_params& params, \ - const raft::device_matrix_view& dataset, \ + const raft::device_matrix_view& dataset, \ raft::neighbors::ivf_pq::index* idx) \ { \ *idx = raft::neighbors::ivf_pq::build(handle, params, dataset); \ } \ void extend(raft::device_resources const& handle, \ raft::neighbors::ivf_pq::index* idx, \ - const raft::device_matrix_view& new_vectors, \ - const raft::device_matrix_view& new_indices) \ + const raft::device_matrix_view& new_vectors, \ + const raft::device_matrix_view& new_indices) \ { \ raft::neighbors::ivf_pq::extend(handle, idx, new_vectors, new_indices); \ } diff --git a/cpp/src/distance/neighbors/ivfpq_search.cu b/cpp/src/distance/neighbors/ivfpq_search.cu index 6ce4cd8375..8c20303149 100644 --- a/cpp/src/distance/neighbors/ivfpq_search.cu +++ b/cpp/src/distance/neighbors/ivfpq_search.cu @@ -20,18 +20,18 @@ namespace raft::runtime::neighbors::ivf_pq { -#define RAFT_SEARCH_INST(T, IdxT) \ - void search(raft::device_resources const& handle, \ - const raft::neighbors::ivf_pq::search_params& params, \ - const raft::neighbors::ivf_pq::index& idx, \ - const raft::device_matrix_view& queries, \ - uint32_t k, \ - const raft::device_matrix_view& neighbors, \ - const raft::device_matrix_view& distances, \ - rmm::mr::device_memory_resource* mr) \ - { \ - raft::neighbors::ivf_pq::search( \ - handle, params, idx, queries, k, neighbors, distances, mr); \ +#define RAFT_SEARCH_INST(T, IdxT) \ + void search(raft::device_resources const& handle, \ + const raft::neighbors::ivf_pq::search_params& params, \ + const raft::neighbors::ivf_pq::index& idx, \ + const raft::device_matrix_view& queries, \ + uint32_t k, \ + const raft::device_matrix_view& neighbors, \ + const raft::device_matrix_view& distances, \ + rmm::mr::device_memory_resource* mr) \ + { \ + raft::neighbors::ivf_pq::search( \ + handle, params, idx, queries, k, neighbors, distances, mr); \ } RAFT_SEARCH_INST(float, uint64_t); From 17ef73c8344836de5b0d93c8a5dba207afdc4100 Mon Sep 17 00:00:00 2001 From: viclafargue Date: Thu, 9 Feb 2023 17:20:37 +0100 Subject: [PATCH 04/20] addressing review --- cpp/include/raft/neighbors/ivf_pq.cuh | 275 +++++++++--------- cpp/include/raft_runtime/neighbors/ivf_pq.hpp | 39 +-- .../neighbors/ivfpq_search_float_uint64_t.cu | 24 +- .../neighbors/ivfpq_search_int8_t_uint64_t.cu | 24 +- .../ivfpq_search_uint8_t_uint64_t.cu | 5 +- .../ivfpq_build_float_uint64_t.cu | 11 +- .../ivfpq_build_int8_t_uint64_t.cu | 10 +- .../ivfpq_build_uint8_t_uint64_t.cu | 10 +- .../ivfpq_extend_float_uint64_t.cu | 24 +- .../ivfpq_extend_int8_t_uint64_t.cu | 24 +- .../ivfpq_extend_uint8_uint64_t.cu | 24 +- python/pylibraft/pylibraft/common/mdspan.pyx | 48 +++ .../neighbors/ivf_pq/cpp/c_ivf_pq.pxd | 52 ++-- .../pylibraft/neighbors/ivf_pq/ivf_pq.pyx | 65 ++--- .../pylibraft/pylibraft/neighbors/refine.pyx | 48 +-- 15 files changed, 309 insertions(+), 374 deletions(-) diff --git a/cpp/include/raft/neighbors/ivf_pq.cuh b/cpp/include/raft/neighbors/ivf_pq.cuh index 4311705303..d392cfd3e5 100644 --- a/cpp/include/raft/neighbors/ivf_pq.cuh +++ b/cpp/include/raft/neighbors/ivf_pq.cuh @@ -32,6 +32,143 @@ namespace raft::neighbors::ivf_pq { * @{ */ +/** + * @brief Build the index from the dataset for efficient search. + * + * NB: Currently, the following distance metrics are supported: + * - L2Expanded + * - L2Unexpanded + * - InnerProduct + * + * @tparam T data element type + * @tparam IdxT type of the indices in the source dataset + * + * @param handle + * @param 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 + */ +template +inline auto build(raft::device_resources const& handle, + const index_params& params, + const raft::device_matrix_view& dataset) -> index +{ + IdxT n_rows = dataset.extent(0); + IdxT dim = dataset.extent(1); + return raft::spatial::knn::ivf_pq::detail::build( + handle, params, dataset.data_handle(), n_rows, dim); +} + +/** + * @brief Build a new index containing the data of the original plus new extra vectors. + * + * Implementation note: + * The new data is clustered according to existing kmeans clusters, then the cluster + * centers are unchanged. + * + * @tparam T data element type + * @tparam IdxT type of the indices in the source dataset + * + * @param handle + * @param orig_index original index + * @param[in] new_vectors a device matrix view to a row-major matrix [n_rows, index.dim()] + * @param[in] new_indices a device matrix view to a vector of indices [n_rows]. + * If the original index is empty (`orig_index.size() == 0`), you can pass `nullptr` + * here to imply a continuous range `[0...n_rows)`. + * + * @return the constructed extended ivf-pq index + */ +template +inline auto extend(raft::device_resources const& handle, + const index& orig_index, + const raft::device_matrix_view& new_vectors, + const raft::device_matrix_view& new_indices) + -> index +{ + IdxT n_rows = new_vectors.extent(0); + ASSERT(n_rows == new_indices.extent(0), + "new_vectors and new_indices have different number of rows"); + ASSERT(new_vectors.extent(1) == orig_index.dim(), + "new_vectors should have the same dimension as the index"); + return raft::spatial::knn::ivf_pq::detail::extend( + handle, orig_index, new_vectors.data_handle(), new_indices.data_handle(), n_rows); +} + +/** + * @brief Extend the index with the new data. + * * + * @tparam T data element type + * @tparam IdxT type of the indices in the source dataset + * + * @param handle + * @param[inout] index + * @param[in] new_vectors a device matrix view to a row-major matrix [n_rows, index.dim()] + * @param[in] new_indices a device matrix view to a vector of indices [n_rows]. + * If the original index is empty (`orig_index.size() == 0`), you can pass `nullptr` + * here to imply a continuous range `[0...n_rows)`. + */ +template +inline void extend(raft::device_resources const& handle, + index* index, + const raft::device_matrix_view& new_vectors, + const raft::device_matrix_view& new_indices) +{ + *index = extend(handle, *index, new_vectors, new_indices); +} + +/** + * @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`. + * + * @tparam T data element type + * @tparam IdxT type of the indices + * + * @param handle + * @param params configure the search + * @param index ivf-pq constructed index + * @param[in] queries a device matrix view to a row-major matrix [n_queries, index->dim()] + * @param k the number of neighbors to find for each query. + * @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] + */ +template +inline void search(raft::device_resources const& handle, + const search_params& params, + const index& index, + const raft::device_matrix_view& queries, + uint32_t k, + const raft::device_matrix_view& neighbors, + const raft::device_matrix_view& distances) +{ + IdxT n_queries = queries.extent(0); + bool check_n_rows = (n_queries == neighbors.extent(0)) && (n_queries == distances.extent(0)); + ASSERT(check_n_rows, + "queries, neighbors and distances parameters have inconsistent number of rows"); + return raft::spatial::knn::ivf_pq::detail::search(handle, + params, + index, + queries.data_handle(), + n_queries, + k, + neighbors.data_handle(), + distances.data_handle(), + handle.get_workspace_resource()); +} + +/** @} */ // end group ivf_pq + /** * @brief Build the index from the dataset for efficient search. * @@ -74,34 +211,6 @@ auto build(raft::device_resources const& handle, return raft::spatial::knn::ivf_pq::detail::build(handle, params, dataset, n_rows, dim); } -/** - * @brief Build the index from the dataset for efficient search. - * - * NB: Currently, the following distance metrics are supported: - * - L2Expanded - * - L2Unexpanded - * - InnerProduct - * - * @tparam T data element type - * @tparam IdxT type of the indices in the source dataset - * - * @param handle - * @param 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 - */ -template -inline auto build(raft::device_resources const& handle, - const index_params& params, - const raft::device_matrix_view& dataset) -> index -{ - IdxT n_rows = dataset.extent(0); - IdxT dim = dataset.extent(1); - return raft::spatial::knn::ivf_pq::detail::build( - handle, params, dataset.data_handle(), n_rows, dim); -} - /** * @brief Build a new index containing the data of the original plus new extra vectors. * @@ -145,39 +254,6 @@ auto extend(raft::device_resources const& handle, handle, orig_index, new_vectors, new_indices, n_rows); } -/** - * @brief Build a new index containing the data of the original plus new extra vectors. - * - * Implementation note: - * The new data is clustered according to existing kmeans clusters, then the cluster - * centers are unchanged. - * - * @tparam T data element type - * @tparam IdxT type of the indices in the source dataset - * - * @param handle - * @param orig_index original index - * @param[in] new_vectors a device matrix view to a row-major matrix [n_rows, index.dim()] - * @param[in] new_indices a device matrix view to a vector of indices [n_rows]. - * If the original index is empty (`orig_index.size() == 0`), you can pass `nullptr` - * here to imply a continuous range `[0...n_rows)`. - * - * @return the constructed extended ivf-pq index - */ -template -inline auto extend(raft::device_resources const& handle, - const index& orig_index, - const raft::device_matrix_view& new_vectors, - const raft::device_matrix_view& new_indices) - -> index -{ - IdxT n_rows = new_vectors.extent(0); - ASSERT(n_rows == new_indices.extent(0), - "new_vectors and new_indices have different number of rows"); - return raft::spatial::knn::ivf_pq::detail::extend( - handle, orig_index, new_vectors.data_handle(), new_indices.data_handle(), n_rows); -} - /** * @brief Extend the index with the new data. * * @@ -202,28 +278,6 @@ void extend(raft::device_resources const& handle, *index = extend(handle, *index, new_vectors, new_indices, n_rows); } -/** - * @brief Extend the index with the new data. - * * - * @tparam T data element type - * @tparam IdxT type of the indices in the source dataset - * - * @param handle - * @param[inout] index - * @param[in] new_vectors a device matrix view to a row-major matrix [n_rows, index.dim()] - * @param[in] new_indices a device matrix view to a vector of indices [n_rows]. - * If the original index is empty (`orig_index.size() == 0`), you can pass `nullptr` - * here to imply a continuous range `[0...n_rows)`. - */ -template -inline void extend(raft::device_resources const& handle, - index* index, - const raft::device_matrix_view& new_vectors, - const raft::device_matrix_view& new_indices) -{ - *index = extend(handle, *index, new_vectors, new_indices); -} - /** * @brief Search ANN using the constructed index. * @@ -281,59 +335,4 @@ void search(raft::device_resources const& handle, handle, params, index, queries, n_queries, k, neighbors, distances, mr); } -/** - * @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`. - * - * @tparam T data element type - * @tparam IdxT type of the indices - * - * @param handle - * @param params configure the search - * @param index ivf-pq constructed index - * @param[in] queries a device matrix view to a row-major matrix [n_queries, index->dim()] - * @param k the number of neighbors to find for each query. - * @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 mr an optional memory resource to use across the searches (you can provide a large enough - * memory pool here to avoid memory allocations within search). - */ -template -inline void search(raft::device_resources const& handle, - const search_params& params, - const index& index, - const raft::device_matrix_view& queries, - uint32_t k, - const raft::device_matrix_view& neighbors, - const raft::device_matrix_view& distances, - rmm::mr::device_memory_resource* mr = nullptr) -{ - IdxT n_queries = queries.extent(0); - bool check_n_rows = (n_queries == neighbors.extent(0)) && (n_queries == distances.extent(0)); - ASSERT(check_n_rows, - "queries, neighbors and distances parameters have inconsistent number of rows"); - return raft::spatial::knn::ivf_pq::detail::search(handle, - params, - index, - queries.data_handle(), - n_queries, - k, - neighbors.data_handle(), - distances.data_handle(), - mr); -} - -/** @} */ // end group ivf_pq - } // namespace raft::neighbors::ivf_pq diff --git a/cpp/include/raft_runtime/neighbors/ivf_pq.hpp b/cpp/include/raft_runtime/neighbors/ivf_pq.hpp index c8981d3f34..3b74f96f93 100644 --- a/cpp/include/raft_runtime/neighbors/ivf_pq.hpp +++ b/cpp/include/raft_runtime/neighbors/ivf_pq.hpp @@ -21,24 +21,13 @@ namespace raft::runtime::neighbors::ivf_pq { #define RAFT_INST_SEARCH(T, IdxT) \ - void search(raft::device_resources const&, \ - const raft::neighbors::ivf_pq::search_params&, \ - const raft::neighbors::ivf_pq::index&, \ - const T*, \ - uint32_t, \ - uint32_t, \ - IdxT*, \ - float*, \ - rmm::mr::device_memory_resource*); \ - \ void search(raft::device_resources const&, \ const raft::neighbors::ivf_pq::search_params&, \ const raft::neighbors::ivf_pq::index&, \ const raft::device_matrix_view&, \ uint32_t, \ const raft::device_matrix_view&, \ - const raft::device_matrix_view&, \ - rmm::mr::device_memory_resource*); + const raft::device_matrix_view&); RAFT_INST_SEARCH(float, uint64_t); RAFT_INST_SEARCH(int8_t, uint64_t); @@ -50,32 +39,6 @@ RAFT_INST_SEARCH(uint8_t, uint64_t); // wrappers, where exception handling is not compatible with return type that has nontrivial // constructor. #define RAFT_INST_BUILD_EXTEND(T, IdxT) \ - auto build(raft::device_resources const& handle, \ - const raft::neighbors::ivf_pq::index_params& params, \ - const T* dataset, \ - IdxT n_rows, \ - uint32_t dim) \ - ->raft::neighbors::ivf_pq::index; \ - \ - auto extend(raft::device_resources const& handle, \ - const raft::neighbors::ivf_pq::index& orig_index, \ - const T* new_vectors, \ - const IdxT* new_indices, \ - IdxT n_rows) \ - ->raft::neighbors::ivf_pq::index; \ - \ - void build(raft::device_resources const& handle, \ - const raft::neighbors::ivf_pq::index_params& params, \ - const T* dataset, \ - IdxT n_rows, \ - uint32_t dim, \ - raft::neighbors::ivf_pq::index* idx); \ - \ - void extend(raft::device_resources const& handle, \ - raft::neighbors::ivf_pq::index* idx, \ - const T* new_vectors, \ - const IdxT* new_indices, \ - IdxT n_rows); \ auto build(raft::device_resources const& handle, \ const raft::neighbors::ivf_pq::index_params& params, \ const raft::device_matrix_view& dataset) \ diff --git a/cpp/src/distance/neighbors/ivfpq_search_float_uint64_t.cu b/cpp/src/distance/neighbors/ivfpq_search_float_uint64_t.cu index c463aa9845..8bba357f04 100644 --- a/cpp/src/distance/neighbors/ivfpq_search_float_uint64_t.cu +++ b/cpp/src/distance/neighbors/ivfpq_search_float_uint64_t.cu @@ -20,19 +20,17 @@ namespace raft::runtime::neighbors::ivf_pq { -#define RAFT_SEARCH_INST(T, IdxT) \ - void search(raft::device_resources const& handle, \ - const raft::neighbors::ivf_pq::search_params& params, \ - const raft::neighbors::ivf_pq::index& idx, \ - const T* queries, \ - uint32_t n_queries, \ - uint32_t k, \ - IdxT* neighbors, \ - float* distances, \ - rmm::mr::device_memory_resource* mr) \ - { \ - raft::neighbors::ivf_pq::search( \ - handle, params, idx, queries, n_queries, k, neighbors, distances, mr); \ +#define RAFT_SEARCH_INST(T, IdxT) \ + void search(raft::device_resources const& handle, \ + const raft::neighbors::ivf_pq::search_params& params, \ + const raft::neighbors::ivf_pq::index& idx, \ + const raft::device_matrix_view& queries, \ + uint32_t k, \ + const raft::device_matrix_view& neighbors, \ + const raft::device_matrix_view& distances) \ + { \ + raft::neighbors::ivf_pq::search( \ + handle, params, idx, queries, k, neighbors, distances); \ } RAFT_SEARCH_INST(float, uint64_t); diff --git a/cpp/src/distance/neighbors/ivfpq_search_int8_t_uint64_t.cu b/cpp/src/distance/neighbors/ivfpq_search_int8_t_uint64_t.cu index ab0dd576b9..4e30595bb4 100644 --- a/cpp/src/distance/neighbors/ivfpq_search_int8_t_uint64_t.cu +++ b/cpp/src/distance/neighbors/ivfpq_search_int8_t_uint64_t.cu @@ -20,19 +20,17 @@ namespace raft::runtime::neighbors::ivf_pq { -#define RAFT_SEARCH_INST(T, IdxT) \ - void search(raft::device_resources const& handle, \ - const raft::neighbors::ivf_pq::search_params& params, \ - const raft::neighbors::ivf_pq::index& idx, \ - const T* queries, \ - uint32_t n_queries, \ - uint32_t k, \ - IdxT* neighbors, \ - float* distances, \ - rmm::mr::device_memory_resource* mr) \ - { \ - raft::neighbors::ivf_pq::search( \ - handle, params, idx, queries, n_queries, k, neighbors, distances, mr); \ +#define RAFT_SEARCH_INST(T, IdxT) \ + void search(raft::device_resources const& handle, \ + const raft::neighbors::ivf_pq::search_params& params, \ + const raft::neighbors::ivf_pq::index& idx, \ + const raft::device_matrix_view& queries, \ + uint32_t k, \ + const raft::device_matrix_view& neighbors, \ + const raft::device_matrix_view& distances) \ + { \ + raft::neighbors::ivf_pq::search( \ + handle, params, idx, queries, k, neighbors, distances); \ } RAFT_SEARCH_INST(int8_t, uint64_t); diff --git a/cpp/src/distance/neighbors/ivfpq_search_uint8_t_uint64_t.cu b/cpp/src/distance/neighbors/ivfpq_search_uint8_t_uint64_t.cu index d9c69ef426..a115b1d1fd 100644 --- a/cpp/src/distance/neighbors/ivfpq_search_uint8_t_uint64_t.cu +++ b/cpp/src/distance/neighbors/ivfpq_search_uint8_t_uint64_t.cu @@ -27,11 +27,10 @@ namespace raft::runtime::neighbors::ivf_pq { const raft::device_matrix_view& queries, \ uint32_t k, \ const raft::device_matrix_view& neighbors, \ - const raft::device_matrix_view& distances, \ - rmm::mr::device_memory_resource* mr) \ + const raft::device_matrix_view& distances) \ { \ raft::neighbors::ivf_pq::search( \ - handle, params, idx, queries, k, neighbors, distances, mr); \ + handle, params, idx, queries, k, neighbors, distances); \ } RAFT_SEARCH_INST(uint8_t, uint64_t); diff --git a/cpp/src/distance/neighbors/specializations/ivfpq_build_float_uint64_t.cu b/cpp/src/distance/neighbors/specializations/ivfpq_build_float_uint64_t.cu index 0831311372..fcb99228f2 100644 --- a/cpp/src/distance/neighbors/specializations/ivfpq_build_float_uint64_t.cu +++ b/cpp/src/distance/neighbors/specializations/ivfpq_build_float_uint64_t.cu @@ -18,13 +18,12 @@ namespace raft::neighbors::ivf_pq { -#define RAFT_INST_BUILD_EXTEND(T, IdxT) \ - template auto build(raft::device_resources const& handle, \ - const index_params& params, \ - const T* dataset, \ - IdxT n_rows, \ - uint32_t dim) \ +#define RAFT_INST_BUILD_EXTEND(T, IdxT) \ + template auto build(raft::device_resources const& handle, \ + const index_params& params, \ + const raft::device_matrix_view& dataset) \ ->index; + RAFT_INST_BUILD_EXTEND(float, uint64_t); #undef RAFT_INST_BUILD_EXTEND diff --git a/cpp/src/distance/neighbors/specializations/ivfpq_build_int8_t_uint64_t.cu b/cpp/src/distance/neighbors/specializations/ivfpq_build_int8_t_uint64_t.cu index 1e2502fe32..addb94556d 100644 --- a/cpp/src/distance/neighbors/specializations/ivfpq_build_int8_t_uint64_t.cu +++ b/cpp/src/distance/neighbors/specializations/ivfpq_build_int8_t_uint64_t.cu @@ -18,12 +18,10 @@ namespace raft::neighbors::ivf_pq { -#define RAFT_INST_BUILD_EXTEND(T, IdxT) \ - template auto build(raft::device_resources const& handle, \ - const index_params& params, \ - const T* dataset, \ - IdxT n_rows, \ - uint32_t dim) \ +#define RAFT_INST_BUILD_EXTEND(T, IdxT) \ + template auto build(raft::device_resources const& handle, \ + const index_params& params, \ + const raft::device_matrix_view& dataset) \ ->index; RAFT_INST_BUILD_EXTEND(int8_t, uint64_t); diff --git a/cpp/src/distance/neighbors/specializations/ivfpq_build_uint8_t_uint64_t.cu b/cpp/src/distance/neighbors/specializations/ivfpq_build_uint8_t_uint64_t.cu index e3336ad95e..0f87a4737a 100644 --- a/cpp/src/distance/neighbors/specializations/ivfpq_build_uint8_t_uint64_t.cu +++ b/cpp/src/distance/neighbors/specializations/ivfpq_build_uint8_t_uint64_t.cu @@ -18,12 +18,10 @@ namespace raft::neighbors::ivf_pq { -#define RAFT_INST_BUILD_EXTEND(T, IdxT) \ - template auto build(raft::device_resources const& handle, \ - const index_params& params, \ - const T* dataset, \ - IdxT n_rows, \ - uint32_t dim) \ +#define RAFT_INST_BUILD_EXTEND(T, IdxT) \ + template auto build(raft::device_resources const& handle, \ + const index_params& params, \ + const raft::device_matrix_view& dataset) \ ->index; RAFT_INST_BUILD_EXTEND(uint8_t, uint64_t); diff --git a/cpp/src/distance/neighbors/specializations/ivfpq_extend_float_uint64_t.cu b/cpp/src/distance/neighbors/specializations/ivfpq_extend_float_uint64_t.cu index 7aa09e3f43..b4d36c9741 100644 --- a/cpp/src/distance/neighbors/specializations/ivfpq_extend_float_uint64_t.cu +++ b/cpp/src/distance/neighbors/specializations/ivfpq_extend_float_uint64_t.cu @@ -18,18 +18,18 @@ namespace raft::neighbors::ivf_pq { -#define RAFT_INST_BUILD_EXTEND(T, IdxT) \ - template auto extend(raft::device_resources const& handle, \ - const index& orig_index, \ - const T* new_vectors, \ - const IdxT* new_indices, \ - IdxT n_rows) \ - ->index; \ - template void extend(raft::device_resources const& handle, \ - index* index, \ - const T* new_vectors, \ - const IdxT* new_indices, \ - IdxT n_rows); +#define RAFT_INST_BUILD_EXTEND(T, IdxT) \ + template auto extend( \ + raft::device_resources const& handle, \ + const index& orig_index, \ + const raft::device_matrix_view& new_vectors, \ + const raft::device_matrix_view& new_indices) \ + ->index; \ + template void extend( \ + raft::device_resources const& handle, \ + index* index, \ + const raft::device_matrix_view& new_vectors, \ + const raft::device_matrix_view& new_indices); RAFT_INST_BUILD_EXTEND(float, uint64_t); diff --git a/cpp/src/distance/neighbors/specializations/ivfpq_extend_int8_t_uint64_t.cu b/cpp/src/distance/neighbors/specializations/ivfpq_extend_int8_t_uint64_t.cu index 440fe6a4a0..2e5568d4d5 100644 --- a/cpp/src/distance/neighbors/specializations/ivfpq_extend_int8_t_uint64_t.cu +++ b/cpp/src/distance/neighbors/specializations/ivfpq_extend_int8_t_uint64_t.cu @@ -18,18 +18,18 @@ namespace raft::neighbors::ivf_pq { -#define RAFT_INST_BUILD_EXTEND(T, IdxT) \ - template auto extend(raft::device_resources const& handle, \ - const index& orig_index, \ - const T* new_vectors, \ - const IdxT* new_indices, \ - IdxT n_rows) \ - ->index; \ - template void extend(raft::device_resources const& handle, \ - index* index, \ - const T* new_vectors, \ - const IdxT* new_indices, \ - IdxT n_rows); +#define RAFT_INST_BUILD_EXTEND(T, IdxT) \ + template auto extend( \ + raft::device_resources const& handle, \ + const index& orig_index, \ + const raft::device_matrix_view& new_vectors, \ + const raft::device_matrix_view& new_indices) \ + ->index; \ + template void extend( \ + raft::device_resources const& handle, \ + index* index, \ + const raft::device_matrix_view& new_vectors, \ + const raft::device_matrix_view& new_indices); RAFT_INST_BUILD_EXTEND(int8_t, uint64_t); diff --git a/cpp/src/distance/neighbors/specializations/ivfpq_extend_uint8_uint64_t.cu b/cpp/src/distance/neighbors/specializations/ivfpq_extend_uint8_uint64_t.cu index 9aee2dc7d1..7add2ad88b 100644 --- a/cpp/src/distance/neighbors/specializations/ivfpq_extend_uint8_uint64_t.cu +++ b/cpp/src/distance/neighbors/specializations/ivfpq_extend_uint8_uint64_t.cu @@ -18,18 +18,18 @@ namespace raft::neighbors::ivf_pq { -#define RAFT_INST_BUILD_EXTEND(T, IdxT) \ - template auto extend(raft::device_resources const& handle, \ - const index& orig_index, \ - const T* new_vectors, \ - const IdxT* new_indices, \ - IdxT n_rows) \ - ->index; \ - template void extend(raft::device_resources const& handle, \ - index* index, \ - const T* new_vectors, \ - const IdxT* new_indices, \ - IdxT n_rows); +#define RAFT_INST_BUILD_EXTEND(T, IdxT) \ + template auto extend( \ + raft::device_resources const& handle, \ + const index& orig_index, \ + const raft::device_matrix_view& new_vectors, \ + const raft::device_matrix_view& new_indices) \ + ->index; \ + template void extend( \ + raft::device_resources const& handle, \ + index* index, \ + const raft::device_matrix_view& new_vectors, \ + const raft::device_matrix_view& new_indices); RAFT_INST_BUILD_EXTEND(uint8_t, uint64_t); diff --git a/python/pylibraft/pylibraft/common/mdspan.pyx b/python/pylibraft/pylibraft/common/mdspan.pyx index ec825495f4..49c72ec16e 100644 --- a/python/pylibraft/pylibraft/common/mdspan.pyx +++ b/python/pylibraft/pylibraft/common/mdspan.pyx @@ -144,3 +144,51 @@ def run_roundtrip_test_for_mdspan(X, fortran_order=False): X2 = np.load(f) assert np.all(X.shape == X2.shape) assert np.all(X == X2) + + +cdef device_matrix_view[float, uint64_t, row_major] \ + get_device_matrix_view_float(array, check_shape=True) except *: + cai = array + if cai.dtype != np.float32: + raise TypeError("dtype %s not supported" % cai.dtype) + if check_shape and len(cai.shape) != 2: + raise ValueError("Expected a 2D array, got %d D" % len(cai.shape)) + shape = (cai.shape[0], cai.shape[1] if len(cai.shape) == 2 else 1) + return make_device_matrix_view[float, uint64_t, row_major]( + cai.data, shape[0], shape[1]) + + +cdef device_matrix_view[uint64_t, uint64_t, row_major] \ + get_device_matrix_view_uint64(array, check_shape=True) except *: + cai = array + if cai.dtype != np.uint64: + raise TypeError("dtype %s not supported" % cai.dtype) + if check_shape and len(cai.shape) != 2: + raise ValueError("Expected a 2D array, got %d D" % len(cai.shape)) + shape = (cai.shape[0], cai.shape[1] if len(cai.shape) == 2 else 1) + return make_device_matrix_view[uint64_t, uint64_t, row_major]( + cai.data, shape[0], shape[1]) + + +cdef device_matrix_view[uint8_t, uint64_t, row_major] \ + get_device_matrix_view_uint8(array, check_shape=True) except *: + cai = array + if cai.dtype != np.uint8: + raise TypeError("dtype %s not supported" % cai.dtype) + if check_shape and len(cai.shape) != 2: + raise ValueError("Expected a 2D array, got %d D" % len(cai.shape)) + shape = (cai.shape[0], cai.shape[1] if len(cai.shape) == 2 else 1) + return make_device_matrix_view[uint8_t, uint64_t, row_major]( + cai.data, shape[0], shape[1]) + + +cdef device_matrix_view[int8_t, uint64_t, row_major] \ + get_device_matrix_view_int8(array, check_shape=True) except *: + cai = array + if cai.dtype != np.int8: + raise TypeError("dtype %s not supported" % cai.dtype) + if check_shape and len(cai.shape) != 2: + raise ValueError("Expected a 2D array, got %d D" % len(cai.shape)) + shape = (cai.shape[0], cai.shape[1] if len(cai.shape) == 2 else 1) + return make_device_matrix_view[int8_t, uint64_t, row_major]( + cai.data, shape[0], shape[1]) diff --git a/python/pylibraft/pylibraft/neighbors/ivf_pq/cpp/c_ivf_pq.pxd b/python/pylibraft/pylibraft/neighbors/ivf_pq/cpp/c_ivf_pq.pxd index c56c3e9d9b..2dbeb8115d 100644 --- a/python/pylibraft/pylibraft/neighbors/ivf_pq/cpp/c_ivf_pq.pxd +++ b/python/pylibraft/pylibraft/neighbors/ivf_pq/cpp/c_ivf_pq.pxd @@ -37,6 +37,7 @@ from libcpp.string cimport string from rmm._lib.memory_resource cimport device_memory_resource from pylibraft.common.handle cimport device_resources +from pylibraft.common.cpp.mdspan cimport device_matrix_view, row_major from pylibraft.distance.distance_type cimport DistanceType @@ -110,72 +111,57 @@ cdef extern from "raft_runtime/neighbors/ivf_pq.hpp" \ cdef void build(const device_resources& handle, const index_params& params, - const float* dataset, - uint64_t n_rows, - uint32_t dim, + const device_matrix_view[float, uint64_t, row_major]& dataset, index[uint64_t]* index) except + cdef void build(const device_resources& handle, const index_params& params, - const int8_t* dataset, - uint64_t n_rows, - uint32_t dim, + const device_matrix_view[int8_t, uint64_t, row_major]& dataset, index[uint64_t]* index) except + cdef void build(const device_resources& handle, const index_params& params, - const uint8_t* dataset, - uint64_t n_rows, - uint32_t dim, + device_matrix_view[uint8_t, uint64_t, row_major]& dataset, index[uint64_t]* index) except + cdef void extend(const device_resources& handle, index[uint64_t]* index, - const float* new_vectors, - const uint64_t* new_indices, - uint64_t n_rows) except + + const device_matrix_view[float, uint64_t, row_major]& new_vectors, + const device_matrix_view[uint64_t, uint64_t, row_major]& new_indices) except + cdef void extend(const device_resources& handle, index[uint64_t]* index, - const int8_t* new_vectors, - const uint64_t* new_indices, - uint64_t n_rows) except + + const device_matrix_view[int8_t, uint64_t, row_major]& new_vectors, + const device_matrix_view[uint64_t, uint64_t, row_major]& new_indices) except + cdef void extend(const device_resources& handle, index[uint64_t]* index, - const uint8_t* new_vectors, - const uint64_t* new_indices, - uint64_t n_rows) except + + const device_matrix_view[uint8_t, uint64_t, row_major]& new_vectors, + const device_matrix_view[uint64_t, uint64_t, row_major]& new_indices) except + cdef void search(const device_resources& handle, const search_params& params, const index[uint64_t]& index, - const float* queries, - uint32_t n_queries, + const device_matrix_view[float, uint64_t, row_major]& queries, uint32_t k, - uint64_t* neighbors, - float* distances, - device_memory_resource* mr) except + + const device_matrix_view[uint64_t, uint64_t, row_major]& neighbors, + const device_matrix_view[float, uint64_t, row_major]& distances) except + cdef void search(const device_resources& handle, const search_params& params, const index[uint64_t]& index, - const int8_t* queries, - uint32_t n_queries, + const device_matrix_view[int8_t, uint64_t, row_major]& queries, uint32_t k, - uint64_t* neighbors, - float* distances, - device_memory_resource* mr) except + + const device_matrix_view[uint64_t, uint64_t, row_major]& neighbors, + const device_matrix_view[float, uint64_t, row_major]& distances) except + cdef void search(const device_resources& handle, const search_params& params, const index[uint64_t]& index, - const uint8_t* queries, - uint32_t n_queries, + const device_matrix_view[uint8_t, uint64_t, row_major]& queries, uint32_t k, - uint64_t* neighbors, - float* distances, - device_memory_resource* mr) except + + const device_matrix_view[uint64_t, uint64_t, row_major]& neighbors, + const device_matrix_view[float, uint64_t, row_major]& distances) except + cdef void serialize(const device_resources& handle, const string& filename, diff --git a/python/pylibraft/pylibraft/neighbors/ivf_pq/ivf_pq.pyx b/python/pylibraft/pylibraft/neighbors/ivf_pq/ivf_pq.pyx index e7b69ddbea..de6368f217 100644 --- a/python/pylibraft/pylibraft/neighbors/ivf_pq/ivf_pq.pyx +++ b/python/pylibraft/pylibraft/neighbors/ivf_pq/ivf_pq.pyx @@ -50,6 +50,15 @@ from pylibraft.common.handle cimport device_resources from pylibraft.common.handle import auto_sync_handle from pylibraft.common.input_validation import is_c_contiguous +from pylibraft.common.cpp.mdspan cimport ( + device_matrix_view, + make_device_matrix_view, + row_major, + get_device_matrix_view_float, + get_device_matrix_view_uint64, + get_device_matrix_view_uint8, + get_device_matrix_view_int8 +) from rmm._lib.memory_resource cimport ( DeviceMemoryResource, @@ -377,7 +386,6 @@ def build(IndexParams index_params, dataset, handle=None): dataset_dt = dataset_cai.dtype _check_input_array(dataset_cai, [np.dtype('float32'), np.dtype('byte'), np.dtype('ubyte')]) - cdef uintptr_t dataset_ptr = dataset_cai.data cdef uint64_t n_rows = dataset_cai.shape[0] cdef uint32_t dim = dataset_cai.shape[1] @@ -393,27 +401,21 @@ def build(IndexParams index_params, dataset, handle=None): with cuda_interruptible(): c_ivf_pq.build(deref(handle_), index_params.params, - dataset_ptr, - n_rows, - dim, + get_device_matrix_view_float(dataset_cai), idx.index) idx.trained = True elif dataset_dt == np.byte: with cuda_interruptible(): c_ivf_pq.build(deref(handle_), index_params.params, - dataset_ptr, - n_rows, - dim, + get_device_matrix_view_int8(dataset_cai), idx.index) idx.trained = True elif dataset_dt == np.ubyte: with cuda_interruptible(): c_ivf_pq.build(deref(handle_), index_params.params, - dataset_ptr, - n_rows, - dim, + get_device_matrix_view_uint8(dataset_cai), idx.index) idx.trained = True else: @@ -505,30 +507,24 @@ def extend(Index index, new_vectors, new_indices, handle=None): if len(idx_cai.shape)!=1: raise ValueError("Indices array is expected to be 1D") - cdef uintptr_t vecs_ptr = vecs_cai.data - cdef uintptr_t idx_ptr = idx_cai.data - if vecs_dt == np.float32: with cuda_interruptible(): c_ivf_pq.extend(deref(handle_), index.index, - vecs_ptr, - idx_ptr, - n_rows) + get_device_matrix_view_float(vecs_cai), + get_device_matrix_view_uint64(idx_cai, check_shape=False)) elif vecs_dt == np.int8: with cuda_interruptible(): c_ivf_pq.extend(deref(handle_), index.index, - vecs_ptr, - idx_ptr, - n_rows) + get_device_matrix_view_int8(vecs_cai), + get_device_matrix_view_uint64(idx_cai, check_shape=False)) elif vecs_dt == np.uint8: with cuda_interruptible(): c_ivf_pq.extend(deref(handle_), index.index, - vecs_ptr, - idx_ptr, - n_rows) + get_device_matrix_view_uint8(vecs_cai), + get_device_matrix_view_uint64(idx_cai, check_shape=False)) else: raise TypeError("query dtype %s not supported" % vecs_dt) @@ -705,7 +701,6 @@ def search(SearchParams search_params, cdef c_ivf_pq.search_params params = search_params.params - cdef uintptr_t queries_ptr = queries_cai.data cdef uintptr_t neighbors_ptr = neighbors_cai.data cdef uintptr_t distances_ptr = distances_cai.data # TODO(tfeher) pass mr_ptr arg @@ -718,34 +713,28 @@ def search(SearchParams search_params, c_ivf_pq.search(deref(handle_), params, deref(index.index), - queries_ptr, - n_queries, + get_device_matrix_view_float(queries_cai), k, - neighbors_ptr, - distances_ptr, - mr_ptr) + get_device_matrix_view_uint64(neighbors_cai), + get_device_matrix_view_float(distances_cai)) elif queries_dt == np.byte: with cuda_interruptible(): c_ivf_pq.search(deref(handle_), params, deref(index.index), - queries_ptr, - n_queries, + get_device_matrix_view_int8(queries_cai), k, - neighbors_ptr, - distances_ptr, - mr_ptr) + get_device_matrix_view_uint64(neighbors_cai), + get_device_matrix_view_float(distances_cai)) elif queries_dt == np.ubyte: with cuda_interruptible(): c_ivf_pq.search(deref(handle_), params, deref(index.index), - queries_ptr, - n_queries, + get_device_matrix_view_uint8(queries_cai), k, - neighbors_ptr, - distances_ptr, - mr_ptr) + get_device_matrix_view_uint64(neighbors_cai), + get_device_matrix_view_float(distances_cai)) else: raise ValueError("query dtype %s not supported" % queries_dt) diff --git a/python/pylibraft/pylibraft/neighbors/refine.pyx b/python/pylibraft/pylibraft/neighbors/refine.pyx index 5c652f7c73..2eca6a8e71 100644 --- a/python/pylibraft/pylibraft/neighbors/refine.pyx +++ b/python/pylibraft/pylibraft/neighbors/refine.pyx @@ -58,6 +58,10 @@ from pylibraft.common.cpp.mdspan cimport ( make_device_matrix_view, make_host_matrix_view, row_major, + get_device_matrix_view_float, + get_device_matrix_view_uint64, + get_device_matrix_view_uint8, + get_device_matrix_view_int8 ) from pylibraft.neighbors.ivf_pq.cpp.c_ivf_pq cimport ( index_params, @@ -125,50 +129,6 @@ cdef extern from "raft_runtime/neighbors/refine.hpp" \ DistanceType metric) except + -cdef device_matrix_view[float, uint64_t, row_major] \ - get_device_matrix_view_float(array) except *: - cai = cai_wrapper(array) - if cai.dtype != np.float32: - raise TypeError("dtype %s not supported" % cai.dtype) - if len(cai.shape) != 2: - raise ValueError("Expected a 2D array, got %d D" % len(cai.shape)) - return make_device_matrix_view[float, uint64_t, row_major]( - cai.data, cai.shape[0], cai.shape[1]) - - -cdef device_matrix_view[uint64_t, uint64_t, row_major] \ - get_device_matrix_view_uint64(array) except *: - cai = cai_wrapper(array) - if cai.dtype != np.uint64: - raise TypeError("dtype %s not supported" % cai.dtype) - if len(cai.shape) != 2: - raise ValueError("Expected a 2D array, got %d D" % len(cai.shape)) - return make_device_matrix_view[uint64_t, uint64_t, row_major]( - cai.data, cai.shape[0], cai.shape[1]) - - -cdef device_matrix_view[uint8_t, uint64_t, row_major] \ - get_device_matrix_view_uint8(array) except *: - cai = cai_wrapper(array) - if cai.dtype != np.uint8: - raise TypeError("dtype %s not supported" % cai.dtype) - if len(cai.shape) != 2: - raise ValueError("Expected a 2D array, got %d D" % len(cai.shape)) - return make_device_matrix_view[uint8_t, uint64_t, row_major]( - cai.data, cai.shape[0], cai.shape[1]) - - -cdef device_matrix_view[int8_t, uint64_t, row_major] \ - get_device_matrix_view_int8(array) except *: - cai = cai_wrapper(array) - if cai.dtype != np.int8: - raise TypeError("dtype %s not supported" % cai.dtype) - if len(cai.shape) != 2: - raise ValueError("Expected a 2D array, got %d D" % len(cai.shape)) - return make_device_matrix_view[int8_t, uint64_t, row_major]( - cai.data, cai.shape[0], cai.shape[1]) - - def _get_array_params(array_interface, check_dtype=None): dtype = np.dtype(array_interface["typestr"]) if check_dtype is None and dtype != check_dtype: From d4e36609567686744cd378eed807e3fb32f260e4 Mon Sep 17 00:00:00 2001 From: viclafargue Date: Fri, 17 Feb 2023 19:10:31 +0100 Subject: [PATCH 05/20] fix style --- .../pylibraft/neighbors/ivf_pq/ivf_pq.pyx | 27 ++++++++++--------- 1 file changed, 15 insertions(+), 12 deletions(-) diff --git a/python/pylibraft/pylibraft/neighbors/ivf_pq/ivf_pq.pyx b/python/pylibraft/pylibraft/neighbors/ivf_pq/ivf_pq.pyx index de6368f217..5b6aabcf88 100644 --- a/python/pylibraft/pylibraft/neighbors/ivf_pq/ivf_pq.pyx +++ b/python/pylibraft/pylibraft/neighbors/ivf_pq/ivf_pq.pyx @@ -50,15 +50,6 @@ from pylibraft.common.handle cimport device_resources from pylibraft.common.handle import auto_sync_handle from pylibraft.common.input_validation import is_c_contiguous -from pylibraft.common.cpp.mdspan cimport ( - device_matrix_view, - make_device_matrix_view, - row_major, - get_device_matrix_view_float, - get_device_matrix_view_uint64, - get_device_matrix_view_uint8, - get_device_matrix_view_int8 -) from rmm._lib.memory_resource cimport ( DeviceMemoryResource, @@ -66,6 +57,15 @@ from rmm._lib.memory_resource cimport ( ) cimport pylibraft.neighbors.ivf_pq.cpp.c_ivf_pq as c_ivf_pq +from pylibraft.common.cpp.mdspan cimport ( + device_matrix_view, + get_device_matrix_view_float, + get_device_matrix_view_int8, + get_device_matrix_view_uint8, + get_device_matrix_view_uint64, + make_device_matrix_view, + row_major, +) from pylibraft.neighbors.ivf_pq.cpp.c_ivf_pq cimport ( index_params, search_params, @@ -512,19 +512,22 @@ def extend(Index index, new_vectors, new_indices, handle=None): c_ivf_pq.extend(deref(handle_), index.index, get_device_matrix_view_float(vecs_cai), - get_device_matrix_view_uint64(idx_cai, check_shape=False)) + get_device_matrix_view_uint64(idx_cai, + check_shape=False)) elif vecs_dt == np.int8: with cuda_interruptible(): c_ivf_pq.extend(deref(handle_), index.index, get_device_matrix_view_int8(vecs_cai), - get_device_matrix_view_uint64(idx_cai, check_shape=False)) + get_device_matrix_view_uint64(idx_cai, + check_shape=False)) elif vecs_dt == np.uint8: with cuda_interruptible(): c_ivf_pq.extend(deref(handle_), index.index, get_device_matrix_view_uint8(vecs_cai), - get_device_matrix_view_uint64(idx_cai, check_shape=False)) + get_device_matrix_view_uint64(idx_cai, + check_shape=False)) else: raise TypeError("query dtype %s not supported" % vecs_dt) From c73cfb2fddf13f8249040a7787180d5f3e15b428 Mon Sep 17 00:00:00 2001 From: viclafargue Date: Wed, 1 Mar 2023 16:09:17 +0100 Subject: [PATCH 06/20] moving helper funcs around --- .../pylibraft/pylibraft/common/cpp/mdspan.pxd | 1 + python/pylibraft/pylibraft/common/mdspan.pxd | 39 +++++++ python/pylibraft/pylibraft/common/mdspan.pyx | 48 ++++++++ .../pylibraft/neighbors/ivf_pq/ivf_pq.pyx | 107 ++++-------------- .../pylibraft/pylibraft/neighbors/refine.pyx | 95 ++++------------ 5 files changed, 135 insertions(+), 155 deletions(-) create mode 100644 python/pylibraft/pylibraft/common/mdspan.pxd diff --git a/python/pylibraft/pylibraft/common/cpp/mdspan.pxd b/python/pylibraft/pylibraft/common/cpp/mdspan.pxd index c3e5abb47e..a8c636f0b7 100644 --- a/python/pylibraft/pylibraft/common/cpp/mdspan.pxd +++ b/python/pylibraft/pylibraft/common/cpp/mdspan.pxd @@ -19,6 +19,7 @@ # cython: embedsignature = True # cython: language_level = 3 +from libc.stdint cimport int8_t, int64_t, uint8_t, uint64_t from libcpp.string cimport string from pylibraft.common.handle cimport device_resources diff --git a/python/pylibraft/pylibraft/common/mdspan.pxd b/python/pylibraft/pylibraft/common/mdspan.pxd new file mode 100644 index 0000000000..2a0bdaca62 --- /dev/null +++ b/python/pylibraft/pylibraft/common/mdspan.pxd @@ -0,0 +1,39 @@ +# +# Copyright (c) 2022-2023, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +# cython: profile=False +# distutils: language = c++ +# cython: embedsignature = True +# cython: language_level = 3 + +from libc.stdint cimport int8_t, uint8_t, uint64_t +from libcpp.string cimport string + +from pylibraft.common.cpp.mdspan cimport device_matrix_view, row_major +from pylibraft.common.handle cimport device_resources + + +cdef device_matrix_view[float, uint64_t, row_major] get_dmv_float( + array, check_shape) except * + +cdef device_matrix_view[uint8_t, uint64_t, row_major] get_dmv_uint8( + array, check_shape) except * + +cdef device_matrix_view[int8_t, uint64_t, row_major] get_dmv_int8( + array, check_shape) except * + +cdef device_matrix_view[uint64_t, uint64_t, row_major] get_dmv_uint64( + array, check_shape) except * diff --git a/python/pylibraft/pylibraft/common/mdspan.pyx b/python/pylibraft/pylibraft/common/mdspan.pyx index fb17a8b1a9..d8524d94b4 100644 --- a/python/pylibraft/pylibraft/common/mdspan.pyx +++ b/python/pylibraft/pylibraft/common/mdspan.pyx @@ -154,3 +154,51 @@ def run_roundtrip_test_for_mdspan(X, fortran_order=False): X2 = np.load(f) assert np.all(X.shape == X2.shape) assert np.all(X == X2) + + +cdef device_matrix_view[float, uint64_t, row_major] \ + get_dmv_float(array, check_shape) except *: + cai = array + if cai.dtype != np.float32: + raise TypeError("dtype %s not supported" % cai.dtype) + if check_shape and len(cai.shape) != 2: + raise ValueError("Expected a 2D array, got %d D" % len(cai.shape)) + shape = (cai.shape[0], cai.shape[1] if len(cai.shape) == 2 else 1) + return make_device_matrix_view[float, uint64_t, row_major]( + cai.data, shape[0], shape[1]) + + +cdef device_matrix_view[uint8_t, uint64_t, row_major] \ + get_dmv_uint8(array, check_shape) except *: + cai = array + if cai.dtype != np.uint8: + raise TypeError("dtype %s not supported" % cai.dtype) + if check_shape and len(cai.shape) != 2: + raise ValueError("Expected a 2D array, got %d D" % len(cai.shape)) + shape = (cai.shape[0], cai.shape[1] if len(cai.shape) == 2 else 1) + return make_device_matrix_view[uint8_t, uint64_t, row_major]( + cai.data, shape[0], shape[1]) + + +cdef device_matrix_view[int8_t, uint64_t, row_major] \ + get_dmv_int8(array, check_shape) except *: + cai = array + if cai.dtype != np.int8: + raise TypeError("dtype %s not supported" % cai.dtype) + if check_shape and len(cai.shape) != 2: + raise ValueError("Expected a 2D array, got %d D" % len(cai.shape)) + shape = (cai.shape[0], cai.shape[1] if len(cai.shape) == 2 else 1) + return make_device_matrix_view[int8_t, uint64_t, row_major]( + cai.data, shape[0], shape[1]) + + +cdef device_matrix_view[uint64_t, uint64_t, row_major] \ + get_dmv_uint64(array, check_shape) except *: + cai = array + if cai.dtype != np.uint64: + raise TypeError("dtype %s not supported" % cai.dtype) + if check_shape and len(cai.shape) != 2: + raise ValueError("Expected a 2D array, got %d D" % len(cai.shape)) + shape = (cai.shape[0], cai.shape[1] if len(cai.shape) == 2 else 1) + return make_device_matrix_view[uint64_t, uint64_t, row_major]( + cai.data, shape[0], shape[1]) diff --git a/python/pylibraft/pylibraft/neighbors/ivf_pq/ivf_pq.pyx b/python/pylibraft/pylibraft/neighbors/ivf_pq/ivf_pq.pyx index c4512715ec..47d8e94e5f 100644 --- a/python/pylibraft/pylibraft/neighbors/ivf_pq/ivf_pq.pyx +++ b/python/pylibraft/pylibraft/neighbors/ivf_pq/ivf_pq.pyx @@ -23,15 +23,7 @@ import warnings import numpy as np from cython.operator cimport dereference as deref -from libc.stdint cimport ( - int8_t, - int32_t, - int64_t, - uint8_t, - uint32_t, - uint64_t, - uintptr_t, -) +from libc.stdint cimport int32_t, int64_t, uint32_t, uint64_t, uintptr_t from libcpp cimport bool, nullptr from libcpp.string cimport string @@ -58,10 +50,12 @@ from rmm._lib.memory_resource cimport ( ) cimport pylibraft.neighbors.ivf_pq.cpp.c_ivf_pq as c_ivf_pq -from pylibraft.common.cpp.mdspan cimport ( - device_matrix_view, - make_device_matrix_view, - row_major, +from pylibraft.common.cpp.mdspan cimport device_matrix_view +from pylibraft.common.mdspan cimport ( + get_dmv_float, + get_dmv_int8, + get_dmv_uint8, + get_dmv_uint64, ) from pylibraft.neighbors.ivf_pq.cpp.c_ivf_pq cimport ( index_params, @@ -69,54 +63,6 @@ from pylibraft.neighbors.ivf_pq.cpp.c_ivf_pq cimport ( ) -cdef device_matrix_view[float, uint64_t, row_major] \ - get_device_matrix_view_float(array, check_shape=True) except *: - cai = array - if cai.dtype != np.float32: - raise TypeError("dtype %s not supported" % cai.dtype) - if check_shape and len(cai.shape) != 2: - raise ValueError("Expected a 2D array, got %d D" % len(cai.shape)) - shape = (cai.shape[0], cai.shape[1] if len(cai.shape) == 2 else 1) - return make_device_matrix_view[float, uint64_t, row_major]( - cai.data, shape[0], shape[1]) - - -cdef device_matrix_view[uint64_t, uint64_t, row_major] \ - get_device_matrix_view_uint64(array, check_shape=True) except *: - cai = array - if cai.dtype != np.uint64: - raise TypeError("dtype %s not supported" % cai.dtype) - if check_shape and len(cai.shape) != 2: - raise ValueError("Expected a 2D array, got %d D" % len(cai.shape)) - shape = (cai.shape[0], cai.shape[1] if len(cai.shape) == 2 else 1) - return make_device_matrix_view[uint64_t, uint64_t, row_major]( - cai.data, shape[0], shape[1]) - - -cdef device_matrix_view[uint8_t, uint64_t, row_major] \ - get_device_matrix_view_uint8(array, check_shape=True) except *: - cai = array - if cai.dtype != np.uint8: - raise TypeError("dtype %s not supported" % cai.dtype) - if check_shape and len(cai.shape) != 2: - raise ValueError("Expected a 2D array, got %d D" % len(cai.shape)) - shape = (cai.shape[0], cai.shape[1] if len(cai.shape) == 2 else 1) - return make_device_matrix_view[uint8_t, uint64_t, row_major]( - cai.data, shape[0], shape[1]) - - -cdef device_matrix_view[int8_t, uint64_t, row_major] \ - get_device_matrix_view_int8(array, check_shape=True) except *: - cai = array - if cai.dtype != np.int8: - raise TypeError("dtype %s not supported" % cai.dtype) - if check_shape and len(cai.shape) != 2: - raise ValueError("Expected a 2D array, got %d D" % len(cai.shape)) - shape = (cai.shape[0], cai.shape[1] if len(cai.shape) == 2 else 1) - return make_device_matrix_view[int8_t, uint64_t, row_major]( - cai.data, shape[0], shape[1]) - - def _get_metric(metric): SUPPORTED_DISTANCES = { "sqeuclidean": DistanceType.L2Expanded, @@ -464,21 +410,21 @@ def build(IndexParams index_params, dataset, handle=None): with cuda_interruptible(): c_ivf_pq.build(deref(handle_), index_params.params, - get_device_matrix_view_float(dataset_cai), + get_dmv_float(dataset_cai, check_shape=True), idx.index) idx.trained = True elif dataset_dt == np.byte: with cuda_interruptible(): c_ivf_pq.build(deref(handle_), index_params.params, - get_device_matrix_view_int8(dataset_cai), + get_dmv_int8(dataset_cai, check_shape=True), idx.index) idx.trained = True elif dataset_dt == np.ubyte: with cuda_interruptible(): c_ivf_pq.build(deref(handle_), index_params.params, - get_device_matrix_view_uint8(dataset_cai), + get_dmv_uint8(dataset_cai, check_shape=True), idx.index) idx.trained = True else: @@ -574,23 +520,20 @@ def extend(Index index, new_vectors, new_indices, handle=None): with cuda_interruptible(): c_ivf_pq.extend(deref(handle_), index.index, - get_device_matrix_view_float(vecs_cai), - get_device_matrix_view_uint64(idx_cai, - check_shape=False)) + get_dmv_float(vecs_cai, check_shape=True), + get_dmv_uint64(idx_cai, check_shape=False)) elif vecs_dt == np.int8: with cuda_interruptible(): c_ivf_pq.extend(deref(handle_), index.index, - get_device_matrix_view_int8(vecs_cai), - get_device_matrix_view_uint64(idx_cai, - check_shape=False)) + get_dmv_int8(vecs_cai, check_shape=True), + get_dmv_uint64(idx_cai, check_shape=False)) elif vecs_dt == np.uint8: with cuda_interruptible(): c_ivf_pq.extend(deref(handle_), index.index, - get_device_matrix_view_uint8(vecs_cai), - get_device_matrix_view_uint64(idx_cai, - check_shape=False)) + get_dmv_uint8(vecs_cai, check_shape=True), + get_dmv_uint64(idx_cai, check_shape=False)) else: raise TypeError("query dtype %s not supported" % vecs_dt) @@ -779,28 +722,28 @@ def search(SearchParams search_params, c_ivf_pq.search(deref(handle_), params, deref(index.index), - get_device_matrix_view_float(queries_cai), + get_dmv_float(queries_cai, check_shape=True), k, - get_device_matrix_view_uint64(neighbors_cai), - get_device_matrix_view_float(distances_cai)) + get_dmv_uint64(neighbors_cai, check_shape=True), + get_dmv_float(distances_cai, check_shape=True)) elif queries_dt == np.byte: with cuda_interruptible(): c_ivf_pq.search(deref(handle_), params, deref(index.index), - get_device_matrix_view_int8(queries_cai), + get_dmv_int8(queries_cai, check_shape=True), k, - get_device_matrix_view_uint64(neighbors_cai), - get_device_matrix_view_float(distances_cai)) + get_dmv_uint64(neighbors_cai, check_shape=True), + get_dmv_float(distances_cai, check_shape=True)) elif queries_dt == np.ubyte: with cuda_interruptible(): c_ivf_pq.search(deref(handle_), params, deref(index.index), - get_device_matrix_view_uint8(queries_cai), + get_dmv_uint8(queries_cai, check_shape=True), k, - get_device_matrix_view_uint64(neighbors_cai), - get_device_matrix_view_float(distances_cai)) + get_dmv_uint64(neighbors_cai, check_shape=True), + get_dmv_float(distances_cai, check_shape=True)) else: raise ValueError("query dtype %s not supported" % queries_dt) diff --git a/python/pylibraft/pylibraft/neighbors/refine.pyx b/python/pylibraft/pylibraft/neighbors/refine.pyx index e5c88b5c33..e01971d681 100644 --- a/python/pylibraft/pylibraft/neighbors/refine.pyx +++ b/python/pylibraft/pylibraft/neighbors/refine.pyx @@ -21,15 +21,7 @@ import numpy as np from cython.operator cimport dereference as deref -from libc.stdint cimport ( - int8_t, - int32_t, - int64_t, - uint8_t, - uint32_t, - uint64_t, - uintptr_t, -) +from libc.stdint cimport int8_t, int64_t, uint8_t, uint64_t, uintptr_t from libcpp cimport bool, nullptr from pylibraft.distance.distance_type cimport DistanceType @@ -56,64 +48,21 @@ cimport pylibraft.neighbors.ivf_pq.cpp.c_ivf_pq as c_ivf_pq from pylibraft.common.cpp.mdspan cimport ( device_matrix_view, host_matrix_view, - make_device_matrix_view, make_host_matrix_view, row_major, ) +from pylibraft.common.mdspan cimport ( + get_dmv_float, + get_dmv_int8, + get_dmv_uint8, + get_dmv_uint64, +) from pylibraft.neighbors.ivf_pq.cpp.c_ivf_pq cimport ( index_params, search_params, ) -cdef device_matrix_view[float, uint64_t, row_major] \ - get_device_matrix_view_float(array, check_shape=True) except *: - cai = array - if cai.dtype != np.float32: - raise TypeError("dtype %s not supported" % cai.dtype) - if check_shape and len(cai.shape) != 2: - raise ValueError("Expected a 2D array, got %d D" % len(cai.shape)) - shape = (cai.shape[0], cai.shape[1] if len(cai.shape) == 2 else 1) - return make_device_matrix_view[float, uint64_t, row_major]( - cai.data, shape[0], shape[1]) - - -cdef device_matrix_view[uint64_t, uint64_t, row_major] \ - get_device_matrix_view_uint64(array, check_shape=True) except *: - cai = array - if cai.dtype != np.uint64: - raise TypeError("dtype %s not supported" % cai.dtype) - if check_shape and len(cai.shape) != 2: - raise ValueError("Expected a 2D array, got %d D" % len(cai.shape)) - shape = (cai.shape[0], cai.shape[1] if len(cai.shape) == 2 else 1) - return make_device_matrix_view[uint64_t, uint64_t, row_major]( - cai.data, shape[0], shape[1]) - - -cdef device_matrix_view[uint8_t, uint64_t, row_major] \ - get_device_matrix_view_uint8(array, check_shape=True) except *: - cai = array - if cai.dtype != np.uint8: - raise TypeError("dtype %s not supported" % cai.dtype) - if check_shape and len(cai.shape) != 2: - raise ValueError("Expected a 2D array, got %d D" % len(cai.shape)) - shape = (cai.shape[0], cai.shape[1] if len(cai.shape) == 2 else 1) - return make_device_matrix_view[uint8_t, uint64_t, row_major]( - cai.data, shape[0], shape[1]) - - -cdef device_matrix_view[int8_t, uint64_t, row_major] \ - get_device_matrix_view_int8(array, check_shape=True) except *: - cai = array - if cai.dtype != np.int8: - raise TypeError("dtype %s not supported" % cai.dtype) - if check_shape and len(cai.shape) != 2: - raise ValueError("Expected a 2D array, got %d D" % len(cai.shape)) - shape = (cai.shape[0], cai.shape[1] if len(cai.shape) == 2 else 1) - return make_device_matrix_view[int8_t, uint64_t, row_major]( - cai.data, shape[0], shape[1]) - - # We omit the const qualifiers in the interface for refine, because cython # has an issue parsing it (https://github.com/cython/cython/issues/4180). cdef extern from "raft_runtime/neighbors/refine.hpp" \ @@ -338,29 +287,29 @@ def _refine_device(dataset, queries, candidates, k, indices, distances, if dataset_cai.dtype == np.float32: with cuda_interruptible(): c_refine(deref(handle_), - get_device_matrix_view_float(dataset), - get_device_matrix_view_float(queries), - get_device_matrix_view_uint64(candidates), - get_device_matrix_view_uint64(indices), - get_device_matrix_view_float(distances), + get_dmv_float(dataset, check_shape=True), + get_dmv_float(queries, check_shape=True), + get_dmv_uint64(candidates, check_shape=True), + get_dmv_uint64(indices, check_shape=True), + get_dmv_float(distances, check_shape=True), c_metric) elif dataset_cai.dtype == np.int8: with cuda_interruptible(): c_refine(deref(handle_), - get_device_matrix_view_int8(dataset), - get_device_matrix_view_int8(queries), - get_device_matrix_view_uint64(candidates), - get_device_matrix_view_uint64(indices), - get_device_matrix_view_float(distances), + get_dmv_int8(dataset, check_shape=True), + get_dmv_int8(queries, check_shape=True), + get_dmv_uint64(candidates, check_shape=True), + get_dmv_uint64(indices, check_shape=True), + get_dmv_float(distances, check_shape=True), c_metric) elif dataset_cai.dtype == np.uint8: with cuda_interruptible(): c_refine(deref(handle_), - get_device_matrix_view_uint8(dataset), - get_device_matrix_view_uint8(queries), - get_device_matrix_view_uint64(candidates), - get_device_matrix_view_uint64(indices), - get_device_matrix_view_float(distances), + get_dmv_uint8(dataset, check_shape=True), + get_dmv_uint8(queries, check_shape=True), + get_dmv_uint64(candidates, check_shape=True), + get_dmv_uint64(indices, check_shape=True), + get_dmv_float(distances, check_shape=True), c_metric) else: raise TypeError("dtype %s not supported" % dataset_cai.dtype) From efebf6fb20e6f36aa420c7c27398275bb4447f20 Mon Sep 17 00:00:00 2001 From: Tamas Bela Feher Date: Tue, 7 Mar 2023 01:24:15 +0100 Subject: [PATCH 07/20] IVF-Flat python wrappers work in progress --- cpp/CMakeLists.txt | 2 + .../raft_runtime/neighbors/ivf_flat.hpp | 84 +++ cpp/src/distance/neighbors/ivf_flat_build.cu | 78 ++ cpp/src/distance/neighbors/ivf_flat_search.cu | 38 + .../pylibraft/neighbors/CMakeLists.txt | 3 +- .../neighbors/ivf_flat/CMakeLists.txt | 24 + .../pylibraft/neighbors/ivf_flat/__init__.pxd | 0 .../pylibraft/neighbors/ivf_flat/__init__.py | 25 + .../neighbors/ivf_flat/cpp/__init__.pxd | 0 .../neighbors/ivf_flat/cpp/__init__.py | 14 + .../neighbors/ivf_flat/cpp/c_ivf_flat.pxd | 150 ++++ .../pylibraft/neighbors/ivf_flat/ivf_flat.pyx | 712 ++++++++++++++++++ 12 files changed, 1129 insertions(+), 1 deletion(-) create mode 100644 cpp/include/raft_runtime/neighbors/ivf_flat.hpp create mode 100644 cpp/src/distance/neighbors/ivf_flat_build.cu create mode 100644 cpp/src/distance/neighbors/ivf_flat_search.cu create mode 100644 python/pylibraft/pylibraft/neighbors/ivf_flat/CMakeLists.txt create mode 100644 python/pylibraft/pylibraft/neighbors/ivf_flat/__init__.pxd create mode 100644 python/pylibraft/pylibraft/neighbors/ivf_flat/__init__.py create mode 100644 python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/__init__.pxd create mode 100644 python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/__init__.py create mode 100644 python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/c_ivf_flat.pxd create mode 100644 python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 7e5b10b227..b925702ce3 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -370,6 +370,8 @@ if(RAFT_COMPILE_DIST_LIBRARY) src/distance/neighbors/ivfpq_search_float_uint64_t.cu src/distance/neighbors/ivfpq_search_int8_t_uint64_t.cu src/distance/neighbors/ivfpq_search_uint8_t_uint64_t.cu + src/distance/neighbors/ivf_flat_search.cu + src/distance/neighbors/ivf_flat_build.cu src/distance/neighbors/specializations/ivfpq_build_float_uint64_t.cu src/distance/neighbors/specializations/ivfpq_build_int8_t_uint64_t.cu src/distance/neighbors/specializations/ivfpq_build_uint8_t_uint64_t.cu diff --git a/cpp/include/raft_runtime/neighbors/ivf_flat.hpp b/cpp/include/raft_runtime/neighbors/ivf_flat.hpp new file mode 100644 index 0000000000..da4d0d4966 --- /dev/null +++ b/cpp/include/raft_runtime/neighbors/ivf_flat.hpp @@ -0,0 +1,84 @@ +/* + * Copyright (c) 2023, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +namespace raft::runtime::neighbors::ivf_flat { + +// We define overloads for build and extend with void return type. This is used in the Cython +// wrappers, where exception handling is not compatible with return type that has nontrivial +// constructor. +#define RAFT_INST_BUILD_EXTEND(T, IdxT) \ + auto build(raft::device_resources const& handle, \ + raft::device_matrix_view dataset, \ + const raft::neighbors::ivf_flat::index_params& params) \ + ->raft::neighbors::ivf_flat::index; \ + \ + auto extend(raft::device_resources const& handle, \ + const raft::neighbors::ivf_flat::index& orig_index, \ + const T* new_vectors, \ + const IdxT* new_indices, \ + IdxT n_rows) \ + ->raft::neighbors::ivf_flat::index; \ + \ + void build(raft::device_resources const& handle, \ + raft::device_matrix_view dataset, \ + const raft::neighbors::ivf_flat::index_params& params, \ + raft::neighbors::ivf_flat::index* idx); \ + \ + void extend(raft::device_resources const& handle, \ + raft::neighbors::ivf_flat::index* idx, \ + const T* new_vectors, \ + const IdxT* new_indices, \ + IdxT n_rows); + +RAFT_INST_BUILD_EXTEND(float, uint64_t) +RAFT_INST_BUILD_EXTEND(int8_t, uint64_t) +RAFT_INST_BUILD_EXTEND(uint8_t, uint64_t) + +#undef RAFT_INST_BUILD_EXTEND + +// /** +// * Save the index to file. +// * +// * Experimental, both the API and the serialization format are subject to change. +// * +// * @param[in] handle the raft handle +// * @param[in] filename the filename for saving the index +// * @param[in] index IVF-PQ index +// * +// */ +// void save(raft::device_resources const& handle, +// const std::string& filename, +// const raft::neighbors::ivf_flat::index& index); + +// /** +// * Load index from file. +// * +// * Experimental, both the API and the serialization format are subject to change. +// * +// * @param[in] handle the raft handle +// * @param[in] filename the name of the file that stores the index +// * @param[in] index IVF-PQ index +// * +// */ +// void load(raft::device_resources const& handle, +// const std::string& filename, +// raft::neighbors::ivf_flat::index* index); + +} // namespace raft::runtime::neighbors::ivf_flat diff --git a/cpp/src/distance/neighbors/ivf_flat_build.cu b/cpp/src/distance/neighbors/ivf_flat_build.cu new file mode 100644 index 0000000000..de0013d35e --- /dev/null +++ b/cpp/src/distance/neighbors/ivf_flat_build.cu @@ -0,0 +1,78 @@ +/* + * Copyright (c) 2023, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +#include + +namespace raft::runtime::neighbors::ivf_flat { + +#define RAFT_INST_BUILD_EXTEND(T, IdxT) \ + auto build(raft::device_resources const& handle, \ + raft::device_matrix_view dataset, \ + const raft::neighbors::ivf_flat::index_params& params) \ + ->raft::neighbors::ivf_flat::index \ + { \ + return raft::neighbors::ivf_flat::build(handle, dataset, params); \ + } \ + auto extend(raft::device_resources const& handle, \ + const raft::neighbors::ivf_flat::index& orig_index, \ + const T* new_vectors, \ + const IdxT* new_indices, \ + IdxT n_rows) \ + ->raft::neighbors::ivf_flat::index \ + { \ + return raft::neighbors::ivf_flat::extend( \ + handle, orig_index, new_vectors, new_indices, n_rows); \ + } \ + \ + void build(raft::device_resources const& handle, \ + raft::device_matrix_view dataset, \ + const raft::neighbors::ivf_flat::index_params& params, \ + raft::neighbors::ivf_flat::index* idx) \ + { \ + *idx = raft::neighbors::ivf_flat::build(handle, dataset, params); \ + } \ + \ + void extend(raft::device_resources const& handle, \ + raft::neighbors::ivf_flat::index* idx, \ + const T* new_vectors, \ + const IdxT* new_indices, \ + IdxT n_rows) \ + { \ + raft::neighbors::ivf_flat::extend(handle, idx, new_vectors, new_indices, n_rows); \ + } + +RAFT_INST_BUILD_EXTEND(float, uint64_t); +RAFT_INST_BUILD_EXTEND(int8_t, uint64_t); +RAFT_INST_BUILD_EXTEND(uint8_t, uint64_t); + +#undef RAFT_INST_BUILD_EXTEND + +// void save(raft::device_resources const& handle, +// const std::string& filename, +// const raft::neighbors::ivf_flat::index& index) +// { +// raft::spatial::knn::ivf_flat::detail::save(handle, filename, index); +// }; + +// void load(raft::device_resources const& handle, +// const std::string& filename, +// raft::neighbors::ivf_flat::index* index) +// { +// if (!index) { RAFT_FAIL("Invalid index pointer"); } +// *index = raft::spatial::knn::ivf_flat::detail::load(handle, filename); +// }; +} // namespace raft::runtime::neighbors::ivf_flat diff --git a/cpp/src/distance/neighbors/ivf_flat_search.cu b/cpp/src/distance/neighbors/ivf_flat_search.cu new file mode 100644 index 0000000000..7e599c4b99 --- /dev/null +++ b/cpp/src/distance/neighbors/ivf_flat_search.cu @@ -0,0 +1,38 @@ +/* + * Copyright (c) 2022-2023, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include +// #include +#include + +namespace raft::runtime::neighbors::ivf_flat { + +#define RAFT_INST_SEARCH(T, IdxT) \ + void search(raft::device_resources const&, \ + const raft::neighbors::ivf_flat::search_params&, \ + const raft::neighbors::ivf_flat::index&, \ + raft::device_matrix_view queries, \ + raft::device_matrix_view neighbors, \ + raft::device_matrix_view distances, \ + uint32_t k); + +RAFT_INST_SEARCH(float, uint64_t); +RAFT_INST_SEARCH(int8_t, uint64_t); +RAFT_INST_SEARCH(uint8_t, uint64_t); + +#undef RAFT_INST_SEARCH + +} // namespace raft::runtime::neighbors::ivf_flat diff --git a/python/pylibraft/pylibraft/neighbors/CMakeLists.txt b/python/pylibraft/pylibraft/neighbors/CMakeLists.txt index ae5fae8201..6fef9288d3 100644 --- a/python/pylibraft/pylibraft/neighbors/CMakeLists.txt +++ b/python/pylibraft/pylibraft/neighbors/CMakeLists.txt @@ -1,5 +1,5 @@ # ============================================================================= -# Copyright (c) 2022, NVIDIA CORPORATION. +# Copyright (c) 2022-2023, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except # in compliance with the License. You may obtain a copy of the License at @@ -23,4 +23,5 @@ rapids_cython_create_modules( LINKED_LIBRARIES "${linked_libraries}" ASSOCIATED_TARGETS raft MODULE_PREFIX neighbors_ ) +add_subdirectory(ivf_flat) add_subdirectory(ivf_pq) diff --git a/python/pylibraft/pylibraft/neighbors/ivf_flat/CMakeLists.txt b/python/pylibraft/pylibraft/neighbors/ivf_flat/CMakeLists.txt new file mode 100644 index 0000000000..f183e17157 --- /dev/null +++ b/python/pylibraft/pylibraft/neighbors/ivf_flat/CMakeLists.txt @@ -0,0 +1,24 @@ +# ============================================================================= +# Copyright (c) 2023, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except +# in compliance with the License. You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software distributed under the License +# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express +# or implied. See the License for the specific language governing permissions and limitations under +# the License. +# ============================================================================= + +# Set the list of Cython files to build +set(cython_sources ivf_flat.pyx) +set(linked_libraries raft::raft raft::distance) + +# Build all of the Cython targets +rapids_cython_create_modules( + CXX + SOURCE_FILES "${cython_sources}" + LINKED_LIBRARIES "${linked_libraries}" ASSOCIATED_TARGETS raft MODULE_PREFIX neighbors_ivfflat_ +) diff --git a/python/pylibraft/pylibraft/neighbors/ivf_flat/__init__.pxd b/python/pylibraft/pylibraft/neighbors/ivf_flat/__init__.pxd new file mode 100644 index 0000000000..e69de29bb2 diff --git a/python/pylibraft/pylibraft/neighbors/ivf_flat/__init__.py b/python/pylibraft/pylibraft/neighbors/ivf_flat/__init__.py new file mode 100644 index 0000000000..58fd88b873 --- /dev/null +++ b/python/pylibraft/pylibraft/neighbors/ivf_flat/__init__.py @@ -0,0 +1,25 @@ +# Copyright (c) 2023, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# + +from .ivf_flat import Index, IndexParams, SearchParams, build, extend, search + +__all__ = [ + "Index", + "IndexParams", + "SearchParams", + "build", + "extend", + "search", +] diff --git a/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/__init__.pxd b/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/__init__.pxd new file mode 100644 index 0000000000..e69de29bb2 diff --git a/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/__init__.py b/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/__init__.py new file mode 100644 index 0000000000..8f2cc34855 --- /dev/null +++ b/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/__init__.py @@ -0,0 +1,14 @@ +# Copyright (c) 2023, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# diff --git a/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/c_ivf_flat.pxd b/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/c_ivf_flat.pxd new file mode 100644 index 0000000000..60a290add3 --- /dev/null +++ b/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/c_ivf_flat.pxd @@ -0,0 +1,150 @@ +# +# Copyright (c) 2023, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# cython: profile=False +# distutils: language = c++ +# cython: embedsignature = True +# cython: language_level = 3 + +import numpy as np + +import pylibraft.common.handle + +from cython.operator cimport dereference as deref +from libc.stdint cimport ( + int8_t, + int64_t, + uint8_t, + uint32_t, + uint64_t, + uintptr_t, +) +from libcpp cimport bool, nullptr +from libcpp.string cimport string + +from rmm._lib.memory_resource cimport device_memory_resource + +from pylibraft.common.cpp.mdspan cimport ( + device_matrix_view, + host_matrix_view, + make_device_matrix_view, + make_host_matrix_view, + row_major, +) +from pylibraft.common.handle cimport device_resources +from pylibraft.distance.distance_type cimport DistanceType +from pylibraft.neighbors.ivf_pq.cpp.c_ivf_pq cimport ( + ann_index, + ann_index_params, + ann_search_params, +) + + +cdef extern from "raft/neighbors/ivf_flat_types.hpp" \ + namespace "raft::neighbors::ivf_flat" nogil: + + cpdef cppclass index_params(ann_index_params): + uint32_t n_lists + uint32_t kmeans_n_iters + double kmeans_trainset_fraction + bool add_data_on_build + bool adaptive_centers + + cdef cppclass index[T, IdxT](ann_index): + index(const device_resources& handle, + DistanceType metric, + uint32_t n_lists, + bool adaptive_centers, + uint32_t dim) + IdxT size() + uint32_t dim() + DistanceType metric() + uint32_t n_lists() + bool adaptive_centers() + + cpdef cppclass search_params(ann_search_params): + uint32_t n_probes + + +cdef extern from "raft_runtime/neighbors/ivf_flat.hpp" \ + namespace "raft::runtime::neighbors::ivf_flat" nogil: + + cdef void build(const device_resources&, + device_matrix_view[float, uint64_t, row_major] dataset, + const index_params& params, + index[float, uint64_t]* index) except + + + cdef void build(const device_resources& handle, + device_matrix_view[uint8_t, uint64_t, row_major] dataset, + const index_params& params, + index[int8_t, uint64_t]* index) except + + + cdef void build(const device_resources& handle, + device_matrix_view[int8_t, uint64_t, row_major] dataset, + const index_params& params, + index[uint8_t, uint64_t]* index) except + + + cdef void extend(const device_resources& handle, + index[float, uint64_t]* index, + const float* new_vectors, + const uint64_t* new_indices, + uint64_t n_rows) except + + + cdef void extend(const device_resources& handle, + index[int8_t, uint64_t]* index, + const int8_t* new_vectors, + const uint64_t* new_indices, + uint64_t n_rows) except + + + cdef void extend(const device_resources& handle, + index[uint8_t, uint64_t]* index, + const uint8_t* new_vectors, + const uint64_t* new_indices, + uint64_t n_rows) except + + + cdef void search( + const device_resources& handle, + const search_params& params, + const index[float, uint64_t]& index, + device_matrix_view[float, uint64_t, row_major] queries, + device_matrix_view[uint64_t, uint64_t, row_major] neighbors, + device_matrix_view[float, uint64_t, row_major] distances, + uint32_t k) except + + + cdef void search( + const device_resources& handle, + const search_params& params, + const index[int8_t, uint64_t]& index, + device_matrix_view[int8_t, uint64_t, row_major] queries, + device_matrix_view[uint64_t, uint64_t, row_major] neighbors, + device_matrix_view[float, uint64_t, row_major] distances, + uint32_t k) except + + + cdef void search( + const device_resources& handle, + const search_params& params, + const index[uint8_t, uint64_t]& index, + device_matrix_view[uint8_t, uint64_t, row_major] queries, + device_matrix_view[uint64_t, uint64_t, row_major] neighbors, + device_matrix_view[float, uint64_t, row_major] distances, + uint32_t k) except + + + # cdef void save(const device_resources& handle, + # const string& filename, + # const index[uint64_t]& index) except + + + # cdef void load(const device_resources& handle, + # const string& filename, + # index[uint64_t]* index) except + diff --git a/python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx b/python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx new file mode 100644 index 0000000000..04bb389818 --- /dev/null +++ b/python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx @@ -0,0 +1,712 @@ +# +# Copyright (c) 2023, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# cython: profile=False +# distutils: language = c++ +# cython: embedsignature = True +# cython: language_level = 3 + +import warnings + +import numpy as np + +from cython.operator cimport dereference as deref +from libc.stdint cimport ( + int8_t, + int64_t, + uint8_t, + uint32_t, + uint64_t, + uintptr_t, +) +from libcpp cimport bool, nullptr +from libcpp.string cimport string + +from pylibraft.distance.distance_type cimport DistanceType + +from pylibraft.common import ( + DeviceResources, + ai_wrapper, + auto_convert_output, + cai_wrapper, + device_ndarray, +) +from pylibraft.common.cai_wrapper import cai_wrapper + +from pylibraft.common.cpp.mdspan cimport device_matrix_view, row_major + +from pylibraft.common.interruptible import cuda_interruptible + +from pylibraft.common.handle cimport device_resources + +from pylibraft.common.handle import auto_sync_handle +from pylibraft.common.input_validation import is_c_contiguous + +from rmm._lib.memory_resource cimport ( + DeviceMemoryResource, + device_memory_resource, +) + +from pylibraft.neighbors.ivf_pq import ( + _check_input_array, + _get_metric, + _get_metric_string, +) + +cimport pylibraft.neighbors.ivf_flat.cpp.c_ivf_flat as c_ivf_flat +from pylibraft.neighbors.ivf_flat.cpp.c_ivf_flat cimport ( + index_params, + search_params, +) + +from pylibraft.neighbors.refine import ( + get_device_matrix_view_float, + get_device_matrix_view_int8, + get_device_matrix_view_uint8, + get_device_matrix_view_uint64, +) + + +cdef class IndexParams: + cdef c_ivf_flat.index_params params + + def __init__(self, *, + n_lists=1024, + metric="sqeuclidean", + kmeans_n_iters=20, + kmeans_trainset_fraction=0.5, + add_data_on_build=True, + bool adaptive_centers=False): + """" + Parameters to build index for IVF-FLAT nearest neighbor search + + Parameters + ---------- + n_list : int, default = 1024 + The number of clusters used in the coarse quantizer. + metric : string denoting the metric type, default="sqeuclidean" + Valid values for metric: ["sqeuclidean", "inner_product", + "euclidean"], where + - sqeuclidean is the euclidean distance without the square root + operation, i.e.: distance(a,b) = \\sum_i (a_i - b_i)^2, + - euclidean is the euclidean distance + - inner product distance is defined as + distance(a, b) = \\sum_i a_i * b_i. + kmeans_n_iters : int, default = 20 + The number of iterations searching for kmeans centers during index + building. + kmeans_trainset_fraction : int, default = 0.5 + If kmeans_trainset_fraction is less than 1, then the dataset is + subsampled, and only n_samples * kmeans_trainset_fraction rows + are used for training. + add_data_on_build : bool, default = True + After training the coarse and fine quantizers, we will populate + the index with the dataset if add_data_on_build == True, otherwise + the index is left empty, and the extend method can be used + to add new vectors to the index. + adaptive_centers : bool, default = False + By default (adaptive_centers = False), the cluster centers are + trained in `ivf_flat::build`, and and never modified in + `ivf_flat::extend`. The alternative behavior (adaptive_centers + = true) is to update the cluster centers for new data when it is + added. In this case, `index.centers()` are always exactly the + centroids of the data in the corresponding clusters. The drawback + of this behavior is that the centroids depend on the order of + adding new data (through the classification of the added data); + that is, `index.centers()` "drift" together with the changing + distribution of the newly added data. + """ + self.params.n_lists = n_lists + self.params.metric = _get_metric(metric) + self.params.metric_arg = 0 + self.params.kmeans_n_iters = kmeans_n_iters + self.params.kmeans_trainset_fraction = kmeans_trainset_fraction + self.params.add_data_on_build = add_data_on_build + self.params.adaptive_centers = adaptive_centers + + @property + def n_lists(self): + return self.params.n_lists + + @property + def metric(self): + return self.params.metric + + @property + def kmeans_n_iters(self): + return self.params.kmeans_n_iters + + @property + def kmeans_trainset_fraction(self): + return self.params.kmeans_trainset_fraction + + @property + def add_data_on_build(self): + return self.params.add_data_on_build + + @property + def adaptive_centers(self): + return self.params.adaptive_centers + + +cdef class Index: + # We store a pointer to the index because it dose not have a trivial + # constructor. + cdef c_ivf_flat.index[float, uint64_t] * index_float + cdef c_ivf_flat.index[int8_t, uint64_t] * index_int8 + cdef c_ivf_flat.index[uint8_t, uint64_t] * index_uint8 + + cdef readonly bool trained + + def __cinit__(self, handle=None): + self.trained = False + if handle is None: + handle = DeviceResources() + cdef device_resources* handle_ = \ + handle.getHandle() + + self.index_int8 = NULL + self.index_uint8 = NULL + # We create a placeholder object. The actual parameter values do + # not matter, it will be replaced with a built index object later. + self.index_float = new c_ivf_flat.index[float, uint64_t]( + deref(handle_), _get_metric("sqeuclidean"), + 1, + False, + 8) + + def __dealloc__(self): + if self.index_float is not NULL: + del self.index_float + if self.index_int8 is not NULL: + del self.index_int8 + if self.index_uint8 is not NULL: + del self.index_uint8 + + def __repr__(self): + m_str = "metric=" + _get_metric_string(self.index.metric()) + attr_str = [ + attr + "=" + str(getattr(self, attr)) + for attr in ["size", "dim", "n_lists", "adaptive_centers"] + ] + attr_str = [m_str] + attr_str + return "Index(type=IVF-FLAT, " + (", ".join(attr_str)) + ")" + + @property + def dim(self): + return self.index[0].dim() + + @property + def size(self): + return self.index[0].size() + + @property + def metric(self): + return self.index[0].metric() + + @property + def n_lists(self): + return self.index[0].n_lists() + + @property + def adaptive_centers(self): + return self.index[0].adaptive_centers() + + +@auto_sync_handle +@auto_convert_output +def build(IndexParams index_params, dataset, handle=None): + """ + Builds an IVF-FLAT index that can be used for nearest neighbor search. + + Parameters + ---------- + index_params : IndexParams object + dataset : CUDA array interface compliant matrix shape (n_samples, dim) + Supported dtype [float, int8, uint8] + {handle_docstring} + + Returns + ------- + index: ivf_flat.Index + + Examples + -------- + + >>> import cupy as cp + + >>> from pylibraft.common import DeviceResources + >>> from pylibraft.neighbors import ivf_flat + + >>> n_samples = 50000 + >>> n_features = 50 + >>> n_queries = 1000 + + >>> dataset = cp.random.random_sample((n_samples, n_features), + ... dtype=cp.float32) + >>> handle = DeviceResources() + >>> index_params = ivf_flat.IndexParams( + ... n_lists=1024, + ... metric="sqeuclidean") + + >>> index = ivf_flat.build(index_params, dataset, handle=handle) + + >>> # Search using the built index + >>> queries = cp.random.random_sample((n_queries, n_features), + ... dtype=cp.float32) + >>> k = 10 + >>> distances, neighbors = ivf_flat.search(ivf_flat.SearchParams(), index, + ... queries, k, handle=handle) + + >>> distances = cp.asarray(distances) + >>> neighbors = cp.asarray(neighbors) + + >>> # pylibraft functions are often asynchronous so the + >>> # handle needs to be explicitly synchronized + >>> handle.sync() + """ + dataset_cai = cai_wrapper(dataset) + dataset_dt = dataset_cai.dtype + _check_input_array(dataset_cai, [np.dtype('float32'), np.dtype('byte'), + np.dtype('ubyte')]) + cdef uintptr_t dataset_ptr = dataset_cai.data + + cdef uint64_t n_rows = dataset_cai.shape[0] + cdef uint32_t dim = dataset_cai.shape[1] + + if handle is None: + handle = DeviceResources() + cdef device_resources* handle_ = \ + handle.getHandle() + + idx = Index() + + cdef device_matrix_view[float, uint64_t, row_major] dataset_v = \ + get_device_matrix_view_float(dataset) + if dataset_dt == np.float32: + with cuda_interruptible(): + c_ivf_flat.build(deref(handle_), + dataset_view, + index_params.params, + idx.index_float) + idx.trained = True + # elif dataset_dt == np.byte: + # with cuda_interruptible(): + # c_ivf_flat.build(deref(handle_), + # get_device_matrix_view_int8(dataset), + # index_params.params, + # idx.index_int8) + # idx.trained = True + # elif dataset_dt == np.ubyte: + # with cuda_interruptible(): + # c_ivf_flat.build(deref(handle_), + # get_device_matrix_view_uint8(dataset), + # index_params.params, + # idx.index_uint8) + # idx.trained = True + # else: + # raise TypeError("dtype %s not supported" % dataset_dt) + + return idx + + +@auto_sync_handle +@auto_convert_output +def extend(Index index, new_vectors, new_indices, handle=None): + """ + Extend an existing index with new vectors. + + Parameters + ---------- + index : ivf_flat.Index + Trained ivf_flat object. + new_vectors : CUDA array interface compliant matrix shape (n_samples, dim) + Supported dtype [float, int8, uint8] + new_indices : CUDA array interface compliant matrix shape (n_samples, dim) + Supported dtype [uint64] + {handle_docstring} + + Returns + ------- + index: ivf_flat.Index + + Examples + -------- + + >>> import cupy as cp + + >>> from pylibraft.common import DeviceResources + >>> from pylibraft.neighbors import ivf_flat + + >>> n_samples = 50000 + >>> n_features = 50 + >>> n_queries = 1000 + + >>> dataset = cp.random.random_sample((n_samples, n_features), + ... dtype=cp.float32) + >>> handle = DeviceResources() + >>> index = ivf_flat.build(ivf_flat.IndexParams(), dataset, handle=handle) + + >>> n_rows = 100 + >>> more_data = cp.random.random_sample((n_rows, n_features), + ... dtype=cp.float32) + >>> indices = index.size + cp.arange(n_rows, dtype=cp.uint64) + >>> index = ivf_flat.extend(index, more_data, indices) + + >>> # Search using the built index + >>> queries = cp.random.random_sample((n_queries, n_features), + ... dtype=cp.float32) + >>> k = 10 + >>> distances, neighbors = ivf_flat.search(ivf_flat.SearchParams(), + ... index, queries, + ... k, handle=handle) + + >>> # pylibraft functions are often asynchronous so the + >>> # handle needs to be explicitly synchronized + >>> handle.sync() + + >>> distances = cp.asarray(distances) + >>> neighbors = cp.asarray(neighbors) + """ + if not index.trained: + raise ValueError("Index need to be built before calling extend.") + + if handle is None: + handle = DeviceResources() + cdef device_resources* handle_ = \ + handle.getHandle() + + vecs_cai = cai_wrapper(new_vectors) + vecs_dt = vecs_cai.dtype + cdef uint64_t n_rows = vecs_cai.shape[0] + cdef uint32_t dim = vecs_cai.shape[1] + + _check_input_array(vecs_cai, [np.dtype('float32'), np.dtype('byte'), + np.dtype('ubyte')], + exp_cols=index.dim) + + idx_cai = cai_wrapper(new_indices) + _check_input_array(idx_cai, [np.dtype('uint64')], exp_rows=n_rows) + if len(idx_cai.shape)!=1: + raise ValueError("Indices array is expected to be 1D") + + cdef uintptr_t vecs_ptr = vecs_cai.data + cdef uintptr_t idx_ptr = idx_cai.data + + if vecs_dt == np.float32: + with cuda_interruptible(): + c_ivf_flat.extend(deref(handle_), + index.index_float, + vecs_ptr, + idx_ptr, + n_rows) + elif vecs_dt == np.int8: + with cuda_interruptible(): + c_ivf_flat.extend(deref(handle_), + index.index_int8, + vecs_ptr, + idx_ptr, + n_rows) + elif vecs_dt == np.uint8: + with cuda_interruptible(): + c_ivf_flat.extend(deref(handle_), + index.index_uint8, + vecs_ptr, + idx_ptr, + n_rows) + else: + raise TypeError("query dtype %s not supported" % vecs_dt) + + return index + + +cdef class SearchParams: + cdef c_ivf_flat.search_params params + + def __init__(self, *, n_probes=20, + lut_dtype=np.float32, + internal_distance_dtype=np.float32): + """ + IVF-FLAT search parameters + + Parameters + ---------- + n_probes: int, default = 1024 + The number of course clusters to select for the fine search. + """ + self.params.n_probes = n_probes + + def __repr__(self): + attr_str = [attr + "=" + str(getattr(self, attr)) + for attr in ["n_probes"]] + return "SearchParams(type=IVF-FLAT, " + (", ".join(attr_str)) + ")" + + @property + def n_probes(self): + return self.params.n_probes + + +@auto_sync_handle +@auto_convert_output +def search(SearchParams search_params, + Index index, + queries, + k, + neighbors=None, + distances=None, + handle=None): + """ + Find the k nearest neighbors for each query. + + Parameters + ---------- + search_params : SearchParams + index : Index + Trained IVF-FLAT index. + queries : CUDA array interface compliant matrix shape (n_samples, dim) + Supported dtype [float, int8, uint8] + k : int + The number of neighbors. + neighbors : Optional CUDA array interface compliant matrix shape + (n_queries, k), dtype 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) + {handle_docstring} + + Examples + -------- + >>> import cupy as cp + + >>> from pylibraft.common import DeviceResources + >>> from pylibraft.neighbors import ivf_flat + + >>> n_samples = 50000 + >>> n_features = 50 + >>> n_queries = 1000 + >>> dataset = cp.random.random_sample((n_samples, n_features), + ... dtype=cp.float32) + + >>> # Build index + >>> handle = DeviceResources() + >>> index = ivf_flat.build(ivf_flat.IndexParams(), dataset, handle=handle) + + >>> # Search using the built index + >>> queries = cp.random.random_sample((n_queries, n_features), + ... dtype=cp.float32) + >>> k = 10 + >>> search_params = ivf_flat.SearchParams( + ... n_probes=20, + ... lut_dtype=cp.float16, + ... internal_distance_dtype=cp.float32 + ... ) + + # TODO update example to set default pool allocator + # (instead of passing an mr) + + >>> # Using a pooling allocator reduces overhead of temporary array + >>> # creation during search. This is useful if multiple searches + >>> # are performad with same query size. + >>> import rmm + >>> mr = rmm.mr.PoolMemoryResource( + ... rmm.mr.CudaMemoryResource(), + ... initial_pool_size=2**29, + ... maximum_pool_size=2**31 + ... ) + >>> distances, neighbors = ivf_flat.search(search_params, index, queries, + ... k, memory_resource=mr, + ... handle=handle) + + >>> # pylibraft functions are often asynchronous so the + >>> # handle needs to be explicitly synchronized + >>> handle.sync() + + >>> neighbors = cp.asarray(neighbors) + >>> distances = cp.asarray(distances) + """ + + if not index.trained: + raise ValueError("Index need to be built before calling search.") + + if handle is None: + handle = DeviceResources() + cdef device_resources* handle_ = \ + handle.getHandle() + + queries_cai = cai_wrapper(queries) + queries_dt = queries_cai.dtype + cdef uint32_t n_queries = queries_cai.shape[0] + + _check_input_array(queries_cai, [np.dtype('float32'), np.dtype('byte'), + np.dtype('ubyte')], + exp_cols=index.dim) + + if neighbors is None: + neighbors = device_ndarray.empty((n_queries, k), dtype='uint64') + + neighbors_cai = cai_wrapper(neighbors) + _check_input_array(neighbors_cai, [np.dtype('uint64')], + exp_rows=n_queries, exp_cols=k) + + if distances is None: + distances = device_ndarray.empty((n_queries, k), dtype='float32') + + distances_cai = cai_wrapper(distances) + _check_input_array(distances_cai, [np.dtype('float32')], + exp_rows=n_queries, exp_cols=k) + + cdef c_ivf_flat.search_params params = search_params.params + + # if queries_dt == np.float32: + # with cuda_interruptible(): + # c_ivf_flat.search(deref(handle_), + # params, + # deref(index.index_float), + # get_device_matrix_view_float(queries), + # get_device_matrix_view_uint64(neighbors), + # get_device_matrix_view_float(distances), + # k) + # elif queries_dt == np.byte: + # with cuda_interruptible(): + # c_ivf_flat.search(deref(handle_), + # params, + # deref(index.index_int8), + # get_device_matrix_view_int8(queries), + # get_device_matrix_view_uint64(neighbors), + # get_device_matrix_view_float(distances), + # k) + # elif queries_dt == np.ubyte: + # with cuda_interruptible(): + # c_ivf_flat.search(deref(handle_), + # params, + # deref(index.index_uint8), + # get_device_matrix_view_uint8(queries), + # get_device_matrix_view_uint64(neighbors), + # get_device_matrix_view_float(distances), + # k) + # else: + # raise ValueError("query dtype %s not supported" % queries_dt) + + return (distances, neighbors) + + +# @auto_sync_handle +# def save(filename, Index index, handle=None): +# """ +# Saves the index to file. + +# Saving / loading the index is experimental. The serialization format is +# subject to change. + +# Parameters +# ---------- +# filename : string +# Name of the file. +# index : Index +# Trained IVF-PQ index. +# {handle_docstring} + +# Examples +# -------- +# >>> import cupy as cp + +# >>> from pylibraft.common import DeviceResources +# >>> from pylibraft.neighbors import ivf_flat + +# >>> n_samples = 50000 +# >>> n_features = 50 +# >>> dataset = cp.random.random_sample((n_samples, n_features), +# ... dtype=cp.float32) + +# >>> # Build index +# >>> handle = DeviceResources() +# >>> index = ivf_flat.build(ivf_flat.IndexParams(), dataset, +# handle=handle) +# >>> ivf_flat.save("my_index.bin", index, handle=handle) +# """ +# if not index.trained: +# raise ValueError("Index need to be built before saving it.") + +# if handle is None: +# handle = DeviceResources() +# cdef device_resources* handle_ = \ +# handle.getHandle() + +# cdef string c_filename = filename.encode('utf-8') + +# c_ivf_flat.save(deref(handle_), c_filename, deref(index.index)) + + +# @auto_sync_handle +# def load(filename, handle=None): +# """ +# Loads index from file. + +# 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. +# {handle_docstring} + +# Returns +# ------- +# index : Index + +# Examples +# -------- +# >>> import cupy as cp + +# >>> from pylibraft.common import DeviceResources +# >>> from pylibraft.neighbors import ivf_flat + +# >>> n_samples = 50000 +# >>> n_features = 50 +# >>> dataset = cp.random.random_sample((n_samples, n_features), +# ... dtype=cp.float32) + +# >>> # Build and save index +# >>> handle = DeviceResources() +# >>> index = ivf_flat.build(ivf_flat.IndexParams(), dataset, +# handle=handle) +# >>> ivf_flat.save("my_index.bin", index, handle=handle) +# >>> del index + +# >>> n_queries = 100 +# >>> queries = cp.random.random_sample((n_queries, n_features), +# ... dtype=cp.float32) +# >>> handle = DeviceResources() +# >>> index = ivf_flat.load("my_index.bin", handle=handle) + +# >>> distances, neighbors = ivf_flat.search( +# ivf_flat.SearchParams(), index, +# ... queries, k=10, handle=handle) +# """ +# if handle is None: +# handle = DeviceResources() +# cdef device_resources* handle_ = \ +# handle.getHandle() + +# cdef string c_filename = filename.encode('utf-8') +# index = Index() + +# c_ivf_flat.load(deref(handle_), c_filename, index.index) +# index.trained = True + +# return index From f3f3cf750c1c11d385696e944a3163336e0323d8 Mon Sep 17 00:00:00 2001 From: viclafargue Date: Tue, 7 Mar 2023 11:53:30 +0100 Subject: [PATCH 08/20] fix refine --- python/pylibraft/pylibraft/common/mdspan.pyx | 12 ++---- .../pylibraft/pylibraft/neighbors/refine.pyx | 38 ++++++++++--------- 2 files changed, 25 insertions(+), 25 deletions(-) diff --git a/python/pylibraft/pylibraft/common/mdspan.pyx b/python/pylibraft/pylibraft/common/mdspan.pyx index d8524d94b4..22afda043d 100644 --- a/python/pylibraft/pylibraft/common/mdspan.pyx +++ b/python/pylibraft/pylibraft/common/mdspan.pyx @@ -157,8 +157,7 @@ def run_roundtrip_test_for_mdspan(X, fortran_order=False): cdef device_matrix_view[float, uint64_t, row_major] \ - get_dmv_float(array, check_shape) except *: - cai = array + get_dmv_float(cai, check_shape) except *: if cai.dtype != np.float32: raise TypeError("dtype %s not supported" % cai.dtype) if check_shape and len(cai.shape) != 2: @@ -169,8 +168,7 @@ cdef device_matrix_view[float, uint64_t, row_major] \ cdef device_matrix_view[uint8_t, uint64_t, row_major] \ - get_dmv_uint8(array, check_shape) except *: - cai = array + get_dmv_uint8(cai, check_shape) except *: if cai.dtype != np.uint8: raise TypeError("dtype %s not supported" % cai.dtype) if check_shape and len(cai.shape) != 2: @@ -181,8 +179,7 @@ cdef device_matrix_view[uint8_t, uint64_t, row_major] \ cdef device_matrix_view[int8_t, uint64_t, row_major] \ - get_dmv_int8(array, check_shape) except *: - cai = array + get_dmv_int8(cai, check_shape) except *: if cai.dtype != np.int8: raise TypeError("dtype %s not supported" % cai.dtype) if check_shape and len(cai.shape) != 2: @@ -193,8 +190,7 @@ cdef device_matrix_view[int8_t, uint64_t, row_major] \ cdef device_matrix_view[uint64_t, uint64_t, row_major] \ - get_dmv_uint64(array, check_shape) except *: - cai = array + get_dmv_uint64(cai, check_shape) except *: if cai.dtype != np.uint64: raise TypeError("dtype %s not supported" % cai.dtype) if check_shape and len(cai.shape) != 2: diff --git a/python/pylibraft/pylibraft/neighbors/refine.pyx b/python/pylibraft/pylibraft/neighbors/refine.pyx index e01971d681..ddc6f115a3 100644 --- a/python/pylibraft/pylibraft/neighbors/refine.pyx +++ b/python/pylibraft/pylibraft/neighbors/refine.pyx @@ -272,6 +272,9 @@ def _refine_device(dataset, queries, candidates, k, indices, distances, raise ValueError("Argument k must be specified if both indices " "and distances arg is None") + queries_cai = cai_wrapper(queries) + dataset_cai = cai_wrapper(dataset) + candidates_cai = cai_wrapper(candidates) n_queries = cai_wrapper(queries).shape[0] if indices is None: @@ -280,36 +283,37 @@ def _refine_device(dataset, queries, candidates, k, indices, distances, if distances is None: distances = device_ndarray.empty((n_queries, k), dtype='float32') - cdef DistanceType c_metric = _get_metric(metric) + indices_cai = cai_wrapper(indices) + distances_cai = cai_wrapper(distances) - dataset_cai = cai_wrapper(dataset) + cdef DistanceType c_metric = _get_metric(metric) if dataset_cai.dtype == np.float32: with cuda_interruptible(): c_refine(deref(handle_), - get_dmv_float(dataset, check_shape=True), - get_dmv_float(queries, check_shape=True), - get_dmv_uint64(candidates, check_shape=True), - get_dmv_uint64(indices, check_shape=True), - get_dmv_float(distances, check_shape=True), + get_dmv_float(dataset_cai, check_shape=True), + get_dmv_float(queries_cai, check_shape=True), + get_dmv_uint64(candidates_cai, check_shape=True), + get_dmv_uint64(indices_cai, check_shape=True), + get_dmv_float(distances_cai, check_shape=True), c_metric) elif dataset_cai.dtype == np.int8: with cuda_interruptible(): c_refine(deref(handle_), - get_dmv_int8(dataset, check_shape=True), - get_dmv_int8(queries, check_shape=True), - get_dmv_uint64(candidates, check_shape=True), - get_dmv_uint64(indices, check_shape=True), - get_dmv_float(distances, check_shape=True), + get_dmv_int8(dataset_cai, check_shape=True), + get_dmv_int8(queries_cai, check_shape=True), + get_dmv_uint64(candidates_cai, check_shape=True), + get_dmv_uint64(indices_cai, check_shape=True), + get_dmv_float(distances_cai, check_shape=True), c_metric) elif dataset_cai.dtype == np.uint8: with cuda_interruptible(): c_refine(deref(handle_), - get_dmv_uint8(dataset, check_shape=True), - get_dmv_uint8(queries, check_shape=True), - get_dmv_uint64(candidates, check_shape=True), - get_dmv_uint64(indices, check_shape=True), - get_dmv_float(distances, check_shape=True), + get_dmv_uint8(dataset_cai, check_shape=True), + get_dmv_uint8(queries_cai, check_shape=True), + get_dmv_uint64(candidates_cai, check_shape=True), + get_dmv_uint64(indices_cai, check_shape=True), + get_dmv_float(distances_cai, check_shape=True), c_metric) else: raise TypeError("dtype %s not supported" % dataset_cai.dtype) From 0815f74b985101fba06da3a1651d3b3133b5d2fe Mon Sep 17 00:00:00 2001 From: divyegala Date: Wed, 8 Mar 2023 10:29:32 -0800 Subject: [PATCH 09/20] add mdspan based apis to runtime --- .../raft_runtime/neighbors/ivf_flat.hpp | 59 ++++++++++------- cpp/src/distance/neighbors/ivf_flat_build.cu | 66 +++++++++---------- cpp/src/distance/neighbors/ivf_flat_search.cu | 16 +++-- 3 files changed, 78 insertions(+), 63 deletions(-) diff --git a/cpp/include/raft_runtime/neighbors/ivf_flat.hpp b/cpp/include/raft_runtime/neighbors/ivf_flat.hpp index da4d0d4966..989bfcd4c6 100644 --- a/cpp/include/raft_runtime/neighbors/ivf_flat.hpp +++ b/cpp/include/raft_runtime/neighbors/ivf_flat.hpp @@ -23,29 +23,27 @@ namespace raft::runtime::neighbors::ivf_flat { // We define overloads for build and extend with void return type. This is used in the Cython // wrappers, where exception handling is not compatible with return type that has nontrivial // constructor. -#define RAFT_INST_BUILD_EXTEND(T, IdxT) \ - auto build(raft::device_resources const& handle, \ - raft::device_matrix_view dataset, \ - const raft::neighbors::ivf_flat::index_params& params) \ - ->raft::neighbors::ivf_flat::index; \ - \ - auto extend(raft::device_resources const& handle, \ - const raft::neighbors::ivf_flat::index& orig_index, \ - const T* new_vectors, \ - const IdxT* new_indices, \ - IdxT n_rows) \ - ->raft::neighbors::ivf_flat::index; \ - \ - void build(raft::device_resources const& handle, \ - raft::device_matrix_view dataset, \ - const raft::neighbors::ivf_flat::index_params& params, \ - raft::neighbors::ivf_flat::index* idx); \ - \ - void extend(raft::device_resources const& handle, \ - raft::neighbors::ivf_flat::index* idx, \ - const T* new_vectors, \ - const IdxT* new_indices, \ - IdxT n_rows); +#define RAFT_INST_BUILD_EXTEND(T, IdxT) \ + auto build(raft::device_resources const& handle, \ + raft::device_matrix_view dataset, \ + const raft::neighbors::ivf_flat::index_params& params) \ + ->raft::neighbors::ivf_flat::index; \ + \ + auto extend(raft::device_resources const& handle, \ + const raft::neighbors::ivf_flat::index& orig_index, \ + raft::device_matrix_view new_vectors, \ + std::optional> new_indices) \ + ->raft::neighbors::ivf_flat::index; \ + \ + void build(raft::device_resources const& handle, \ + raft::device_matrix_view dataset, \ + const raft::neighbors::ivf_flat::index_params& params, \ + raft::neighbors::ivf_flat::index* idx); \ + \ + void extend(raft::device_resources const& handle, \ + raft::neighbors::ivf_flat::index* idx, \ + raft::device_matrix_view new_vectors, \ + std::optional> new_indices); RAFT_INST_BUILD_EXTEND(float, uint64_t) RAFT_INST_BUILD_EXTEND(int8_t, uint64_t) @@ -53,6 +51,21 @@ RAFT_INST_BUILD_EXTEND(uint8_t, uint64_t) #undef RAFT_INST_BUILD_EXTEND +#define RAFT_INST_SEARCH(T, IdxT) \ + void search(raft::device_resources const&, \ + const raft::neighbors::ivf_flat::index&, \ + raft::device_matrix_view, \ + raft::device_matrix_view, \ + raft::device_matrix_view, \ + raft::neighbors::ivf_flat::search_params const&, \ + uint32_t); + +RAFT_INST_SEARCH(float, uint64_t); +RAFT_INST_SEARCH(int8_t, uint64_t); +RAFT_INST_SEARCH(uint8_t, uint64_t); + +#undef RAFT_INST_SEARCH + // /** // * Save the index to file. // * diff --git a/cpp/src/distance/neighbors/ivf_flat_build.cu b/cpp/src/distance/neighbors/ivf_flat_build.cu index de0013d35e..e31e81e070 100644 --- a/cpp/src/distance/neighbors/ivf_flat_build.cu +++ b/cpp/src/distance/neighbors/ivf_flat_build.cu @@ -19,40 +19,38 @@ namespace raft::runtime::neighbors::ivf_flat { -#define RAFT_INST_BUILD_EXTEND(T, IdxT) \ - auto build(raft::device_resources const& handle, \ - raft::device_matrix_view dataset, \ - const raft::neighbors::ivf_flat::index_params& params) \ - ->raft::neighbors::ivf_flat::index \ - { \ - return raft::neighbors::ivf_flat::build(handle, dataset, params); \ - } \ - auto extend(raft::device_resources const& handle, \ - const raft::neighbors::ivf_flat::index& orig_index, \ - const T* new_vectors, \ - const IdxT* new_indices, \ - IdxT n_rows) \ - ->raft::neighbors::ivf_flat::index \ - { \ - return raft::neighbors::ivf_flat::extend( \ - handle, orig_index, new_vectors, new_indices, n_rows); \ - } \ - \ - void build(raft::device_resources const& handle, \ - raft::device_matrix_view dataset, \ - const raft::neighbors::ivf_flat::index_params& params, \ - raft::neighbors::ivf_flat::index* idx) \ - { \ - *idx = raft::neighbors::ivf_flat::build(handle, dataset, params); \ - } \ - \ - void extend(raft::device_resources const& handle, \ - raft::neighbors::ivf_flat::index* idx, \ - const T* new_vectors, \ - const IdxT* new_indices, \ - IdxT n_rows) \ - { \ - raft::neighbors::ivf_flat::extend(handle, idx, new_vectors, new_indices, n_rows); \ +#define RAFT_INST_BUILD_EXTEND(T, IdxT) \ + auto build(raft::device_resources const& handle, \ + raft::device_matrix_view dataset, \ + const raft::neighbors::ivf_flat::index_params& params) \ + ->raft::neighbors::ivf_flat::index \ + { \ + return raft::neighbors::ivf_flat::build(handle, dataset, params); \ + } \ + auto extend(raft::device_resources const& handle, \ + const raft::neighbors::ivf_flat::index& orig_index, \ + raft::device_matrix_view new_vectors, \ + std::optional> new_indices) \ + ->raft::neighbors::ivf_flat::index \ + { \ + return raft::neighbors::ivf_flat::extend( \ + handle, orig_index, new_vectors, new_indices); \ + } \ + \ + void build(raft::device_resources const& handle, \ + raft::device_matrix_view dataset, \ + const raft::neighbors::ivf_flat::index_params& params, \ + raft::neighbors::ivf_flat::index* idx) \ + { \ + *idx = raft::neighbors::ivf_flat::build(handle, dataset, params); \ + } \ + \ + void extend(raft::device_resources const& handle, \ + raft::neighbors::ivf_flat::index* idx, \ + raft::device_matrix_view new_vectors, \ + std::optional> new_indices) \ + { \ + raft::neighbors::ivf_flat::extend(handle, idx, new_vectors, new_indices); \ } RAFT_INST_BUILD_EXTEND(float, uint64_t); diff --git a/cpp/src/distance/neighbors/ivf_flat_search.cu b/cpp/src/distance/neighbors/ivf_flat_search.cu index 7e599c4b99..50441e25d8 100644 --- a/cpp/src/distance/neighbors/ivf_flat_search.cu +++ b/cpp/src/distance/neighbors/ivf_flat_search.cu @@ -14,20 +14,24 @@ * limitations under the License. */ -#include +#include // #include #include namespace raft::runtime::neighbors::ivf_flat { #define RAFT_INST_SEARCH(T, IdxT) \ - void search(raft::device_resources const&, \ - const raft::neighbors::ivf_flat::search_params&, \ - const raft::neighbors::ivf_flat::index&, \ + void search(raft::device_resources const& handle, \ + const raft::neighbors::ivf_flat::index& index, \ raft::device_matrix_view queries, \ raft::device_matrix_view neighbors, \ - raft::device_matrix_view distances, \ - uint32_t k); + raft::device_matrix_view distances, \ + raft::neighbors::ivf_flat::search_params const& params, \ + uint32_t k) \ + { \ + raft::neighbors::ivf_flat::search( \ + handle, index, queries, neighbors, distances, params, k); \ + } RAFT_INST_SEARCH(float, uint64_t); RAFT_INST_SEARCH(int8_t, uint64_t); From 89c0b33f9620533c59a1f077eb43169352369bb5 Mon Sep 17 00:00:00 2001 From: divyegala Date: Wed, 8 Mar 2023 10:48:52 -0800 Subject: [PATCH 10/20] link to runtime, build successfully --- .../neighbors/ivf_flat/cpp/c_ivf_flat.pxd | 42 ++--- .../pylibraft/neighbors/ivf_flat/ivf_flat.pyx | 163 ++++++++++-------- 2 files changed, 114 insertions(+), 91 deletions(-) diff --git a/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/c_ivf_flat.pxd b/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/c_ivf_flat.pxd index 60a290add3..577612593a 100644 --- a/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/c_ivf_flat.pxd +++ b/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/c_ivf_flat.pxd @@ -38,11 +38,13 @@ from rmm._lib.memory_resource cimport device_memory_resource from pylibraft.common.cpp.mdspan cimport ( device_matrix_view, + device_vector_view, host_matrix_view, make_device_matrix_view, make_host_matrix_view, row_major, ) +from pylibraft.common.cpp.optional cimport optional from pylibraft.common.handle cimport device_resources from pylibraft.distance.distance_type cimport DistanceType from pylibraft.neighbors.ivf_pq.cpp.c_ivf_pq cimport ( @@ -87,58 +89,58 @@ cdef extern from "raft_runtime/neighbors/ivf_flat.hpp" \ index[float, uint64_t]* index) except + cdef void build(const device_resources& handle, - device_matrix_view[uint8_t, uint64_t, row_major] dataset, + device_matrix_view[int8_t, uint64_t, row_major] dataset, const index_params& params, index[int8_t, uint64_t]* index) except + cdef void build(const device_resources& handle, - device_matrix_view[int8_t, uint64_t, row_major] dataset, + device_matrix_view[uint8_t, uint64_t, row_major] dataset, const index_params& params, index[uint8_t, uint64_t]* index) except + - cdef void extend(const device_resources& handle, - index[float, uint64_t]* index, - const float* new_vectors, - const uint64_t* new_indices, - uint64_t n_rows) except + + cdef void extend( + const device_resources& handle, + index[float, uint64_t]* index, + device_matrix_view[float, uint64_t, row_major] new_vectors, + optional[device_vector_view[uint64_t, uint64_t]] new_indices) except + - cdef void extend(const device_resources& handle, - index[int8_t, uint64_t]* index, - const int8_t* new_vectors, - const uint64_t* new_indices, - uint64_t n_rows) except + + cdef void extend( + const device_resources& handle, + index[int8_t, uint64_t]* index, + device_matrix_view[int8_t, uint64_t, row_major] new_vectors, + optional[device_vector_view[uint64_t, uint64_t]] new_indices) except + - cdef void extend(const device_resources& handle, - index[uint8_t, uint64_t]* index, - const uint8_t* new_vectors, - const uint64_t* new_indices, - uint64_t n_rows) except + + cdef void extend( + const device_resources& handle, + index[uint8_t, uint64_t]* index, + device_matrix_view[uint8_t, uint64_t, row_major] new_vectors, + optional[device_vector_view[uint64_t, uint64_t]] new_indices) except + cdef void search( const device_resources& handle, - const search_params& params, const index[float, uint64_t]& index, device_matrix_view[float, uint64_t, row_major] queries, device_matrix_view[uint64_t, uint64_t, row_major] neighbors, device_matrix_view[float, uint64_t, row_major] distances, + const search_params& params, uint32_t k) except + cdef void search( const device_resources& handle, - const search_params& params, const index[int8_t, uint64_t]& index, device_matrix_view[int8_t, uint64_t, row_major] queries, device_matrix_view[uint64_t, uint64_t, row_major] neighbors, device_matrix_view[float, uint64_t, row_major] distances, + const search_params& params, uint32_t k) except + cdef void search( const device_resources& handle, - const search_params& params, const index[uint8_t, uint64_t]& index, device_matrix_view[uint8_t, uint64_t, row_major] queries, device_matrix_view[uint64_t, uint64_t, row_major] neighbors, device_matrix_view[float, uint64_t, row_major] distances, + const search_params& params, uint32_t k) except + # cdef void save(const device_resources& handle, diff --git a/python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx b/python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx index 04bb389818..16b259c7c1 100644 --- a/python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx +++ b/python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx @@ -45,7 +45,12 @@ from pylibraft.common import ( ) from pylibraft.common.cai_wrapper import cai_wrapper -from pylibraft.common.cpp.mdspan cimport device_matrix_view, row_major +from pylibraft.common.cpp.mdspan cimport ( + device_matrix_view, + device_vector_view, + make_device_vector_view, + row_major, +) from pylibraft.common.interruptible import cuda_interruptible @@ -59,6 +64,8 @@ from rmm._lib.memory_resource cimport ( device_memory_resource, ) +from pylibraft.common.cpp.optional cimport optional + from pylibraft.neighbors.ivf_pq import ( _check_input_array, _get_metric, @@ -66,18 +73,17 @@ from pylibraft.neighbors.ivf_pq import ( ) cimport pylibraft.neighbors.ivf_flat.cpp.c_ivf_flat as c_ivf_flat +from pylibraft.common.mdspan cimport ( + get_dmv_float, + get_dmv_int8, + get_dmv_uint8, + get_dmv_uint64, +) from pylibraft.neighbors.ivf_flat.cpp.c_ivf_flat cimport ( index_params, search_params, ) -from pylibraft.neighbors.refine import ( - get_device_matrix_view_float, - get_device_matrix_view_int8, - get_device_matrix_view_uint8, - get_device_matrix_view_uint64, -) - cdef class IndexParams: cdef c_ivf_flat.index_params params @@ -177,8 +183,6 @@ cdef class Index: cdef device_resources* handle_ = \ handle.getHandle() - self.index_int8 = NULL - self.index_uint8 = NULL # We create a placeholder object. The actual parameter values do # not matter, it will be replaced with a built index object later. self.index_float = new c_ivf_flat.index[float, uint64_t]( @@ -187,6 +191,18 @@ cdef class Index: False, 8) + self.index_int8 = new c_ivf_flat.index[int8_t, uint64_t]( + deref(handle_), _get_metric("sqeuclidean"), + 1, + False, + 8) + + self.index_uint8 = new c_ivf_flat.index[uint8_t, uint64_t]( + deref(handle_), _get_metric("sqeuclidean"), + 1, + False, + 8) + def __dealloc__(self): if self.index_float is not NULL: del self.index_float @@ -281,7 +297,6 @@ def build(IndexParams index_params, dataset, handle=None): dataset_dt = dataset_cai.dtype _check_input_array(dataset_cai, [np.dtype('float32'), np.dtype('byte'), np.dtype('ubyte')]) - cdef uintptr_t dataset_ptr = dataset_cai.data cdef uint64_t n_rows = dataset_cai.shape[0] cdef uint32_t dim = dataset_cai.shape[1] @@ -293,32 +308,28 @@ def build(IndexParams index_params, dataset, handle=None): idx = Index() - cdef device_matrix_view[float, uint64_t, row_major] dataset_v = \ - get_device_matrix_view_float(dataset) if dataset_dt == np.float32: with cuda_interruptible(): c_ivf_flat.build(deref(handle_), - dataset_view, + get_dmv_float(dataset, check_shape=True), index_params.params, idx.index_float) - idx.trained = True - # elif dataset_dt == np.byte: - # with cuda_interruptible(): - # c_ivf_flat.build(deref(handle_), - # get_device_matrix_view_int8(dataset), - # index_params.params, - # idx.index_int8) - # idx.trained = True - # elif dataset_dt == np.ubyte: - # with cuda_interruptible(): - # c_ivf_flat.build(deref(handle_), - # get_device_matrix_view_uint8(dataset), - # index_params.params, - # idx.index_uint8) - # idx.trained = True - # else: - # raise TypeError("dtype %s not supported" % dataset_dt) + elif dataset_dt == np.byte: + with cuda_interruptible(): + c_ivf_flat.build(deref(handle_), + get_dmv_int8(dataset, check_shape=True), + index_params.params, + idx.index_int8) + elif dataset_dt == np.ubyte: + with cuda_interruptible(): + c_ivf_flat.build(deref(handle_), + get_dmv_uint8(dataset, check_shape=True), + index_params.params, + idx.index_uint8) + else: + raise TypeError("dtype %s not supported" % dataset_dt) + idx.trained = True return idx @@ -402,30 +413,40 @@ def extend(Index index, new_vectors, new_indices, handle=None): if len(idx_cai.shape)!=1: raise ValueError("Indices array is expected to be 1D") - cdef uintptr_t vecs_ptr = vecs_cai.data - cdef uintptr_t idx_ptr = idx_cai.data + cdef optional[device_vector_view[uint64_t, uint64_t]] new_indices_float + cdef optional[device_vector_view[uint64_t, uint64_t]] new_indices_int8 + cdef optional[device_vector_view[uint64_t, uint64_t]] new_indices_uint8 if vecs_dt == np.float32: + if index.index_float[0].size() > 0: + new_indices_float = make_device_vector_view( + idx_cai.data, + idx_cai.shape[0]) with cuda_interruptible(): c_ivf_flat.extend(deref(handle_), index.index_float, - vecs_ptr, - idx_ptr, - n_rows) + get_dmv_float(new_vectors, check_shape=True), + new_indices_float) elif vecs_dt == np.int8: + if index.index_int8[0].size() > 0: + new_indices_int8 = make_device_vector_view( + idx_cai.data, + idx_cai.shape[0]) with cuda_interruptible(): c_ivf_flat.extend(deref(handle_), index.index_int8, - vecs_ptr, - idx_ptr, - n_rows) + get_dmv_int8(new_vectors, check_shape=True), + new_indices_int8) elif vecs_dt == np.uint8: + if index.index_uint8[0].size() > 0: + new_indices_uint8 = make_device_vector_view( + idx_cai.data, + idx_cai.shape[0]) with cuda_interruptible(): c_ivf_flat.extend(deref(handle_), index.index_uint8, - vecs_ptr, - idx_ptr, - n_rows) + get_dmv_uint8(new_vectors, check_shape=True), + new_indices_uint8) else: raise TypeError("query dtype %s not supported" % vecs_dt) @@ -570,35 +591,35 @@ def search(SearchParams search_params, cdef c_ivf_flat.search_params params = search_params.params - # if queries_dt == np.float32: - # with cuda_interruptible(): - # c_ivf_flat.search(deref(handle_), - # params, - # deref(index.index_float), - # get_device_matrix_view_float(queries), - # get_device_matrix_view_uint64(neighbors), - # get_device_matrix_view_float(distances), - # k) - # elif queries_dt == np.byte: - # with cuda_interruptible(): - # c_ivf_flat.search(deref(handle_), - # params, - # deref(index.index_int8), - # get_device_matrix_view_int8(queries), - # get_device_matrix_view_uint64(neighbors), - # get_device_matrix_view_float(distances), - # k) - # elif queries_dt == np.ubyte: - # with cuda_interruptible(): - # c_ivf_flat.search(deref(handle_), - # params, - # deref(index.index_uint8), - # get_device_matrix_view_uint8(queries), - # get_device_matrix_view_uint64(neighbors), - # get_device_matrix_view_float(distances), - # k) - # else: - # raise ValueError("query dtype %s not supported" % queries_dt) + if queries_dt == np.float32: + with cuda_interruptible(): + c_ivf_flat.search(deref(handle_), + deref(index.index_float), + get_dmv_float(queries, check_shape=True), + get_dmv_uint64(neighbors, check_shape=True), + get_dmv_float(distances, check_shape=True), + params, + k) + elif queries_dt == np.byte: + with cuda_interruptible(): + c_ivf_flat.search(deref(handle_), + deref(index.index_int8), + get_dmv_int8(queries, check_shape=True), + get_dmv_uint64(neighbors, check_shape=True), + get_dmv_float(distances, check_shape=True), + params, + k) + elif queries_dt == np.ubyte: + with cuda_interruptible(): + c_ivf_flat.search(deref(handle_), + deref(index.index_uint8), + get_dmv_uint8(queries, check_shape=True), + get_dmv_uint64(neighbors, check_shape=True), + get_dmv_float(distances, check_shape=True), + params, + k) + else: + raise ValueError("query dtype %s not supported" % queries_dt) return (distances, neighbors) From 2fd6a60c51bb3b5b1f9f8360488e659662b2e0a5 Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 9 Mar 2023 06:20:48 -0800 Subject: [PATCH 11/20] add index pointer API for mdspan build, remove k from mdspan search, add specializations --- cpp/CMakeLists.txt | 9 +++ cpp/include/raft/neighbors/ivf_flat.cuh | 61 ++++++++++++++++--- .../raft/neighbors/specializations.cuh | 1 + .../neighbors/specializations/ivf_flat.cuh | 59 ++++++++++++++++++ .../raft_runtime/neighbors/ivf_flat.hpp | 3 +- cpp/src/distance/neighbors/ivf_flat_build.cu | 2 +- cpp/src/distance/neighbors/ivf_flat_search.cu | 10 ++- .../ivfflat_build_float_uint64_t.cu | 36 +++++++++++ .../ivfflat_build_int8_t_uint64_t.cu | 36 +++++++++++ .../ivfflat_build_uint8_t_uint64_t.cu | 36 +++++++++++ .../ivfflat_extend_float_uint64_t.cu | 37 +++++++++++ .../ivfflat_extend_int8_t_uint64_t.cu | 37 +++++++++++ .../ivfflat_extend_uint8_t_uint64_t.cu | 36 +++++++++++ .../ivfflat_search_float_uint64_t.cu | 33 ++++++++++ .../ivfflat_search_int8_t_uint64_t.cu | 33 ++++++++++ .../ivfflat_search_uint8_t_uint64_t.cu | 33 ++++++++++ cpp/test/neighbors/ann_ivf_flat.cuh | 3 +- .../neighbors/ivf_flat/cpp/c_ivf_flat.pxd | 9 +-- .../pylibraft/neighbors/ivf_flat/ivf_flat.pyx | 9 +-- 19 files changed, 451 insertions(+), 32 deletions(-) create mode 100644 cpp/include/raft/neighbors/specializations/ivf_flat.cuh create mode 100644 cpp/src/distance/neighbors/specializations/ivfflat_build_float_uint64_t.cu create mode 100644 cpp/src/distance/neighbors/specializations/ivfflat_build_int8_t_uint64_t.cu create mode 100644 cpp/src/distance/neighbors/specializations/ivfflat_build_uint8_t_uint64_t.cu create mode 100644 cpp/src/distance/neighbors/specializations/ivfflat_extend_float_uint64_t.cu create mode 100644 cpp/src/distance/neighbors/specializations/ivfflat_extend_int8_t_uint64_t.cu create mode 100644 cpp/src/distance/neighbors/specializations/ivfflat_extend_uint8_t_uint64_t.cu create mode 100644 cpp/src/distance/neighbors/specializations/ivfflat_search_float_uint64_t.cu create mode 100644 cpp/src/distance/neighbors/specializations/ivfflat_search_int8_t_uint64_t.cu create mode 100644 cpp/src/distance/neighbors/specializations/ivfflat_search_uint8_t_uint64_t.cu diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index b925702ce3..5bb4f86bd2 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -372,6 +372,15 @@ if(RAFT_COMPILE_DIST_LIBRARY) src/distance/neighbors/ivfpq_search_uint8_t_uint64_t.cu src/distance/neighbors/ivf_flat_search.cu src/distance/neighbors/ivf_flat_build.cu + src/distance/neighbors/specializations/ivfflat_build_float_uint64_t.cu + src/distance/neighbors/specializations/ivfflat_build_int8_t_uint64_t.cu + src/distance/neighbors/specializations/ivfflat_build_uint8_t_uint64_t.cu + src/distance/neighbors/specializations/ivfflat_extend_float_uint64_t.cu + src/distance/neighbors/specializations/ivfflat_extend_int8_t_uint64_t.cu + src/distance/neighbors/specializations/ivfflat_extend_uint8_t_uint64_t.cu + src/distance/neighbors/specializations/ivfflat_search_float_uint64_t.cu + src/distance/neighbors/specializations/ivfflat_search_int8_t_uint64_t.cu + src/distance/neighbors/specializations/ivfflat_search_uint8_t_uint64_t.cu src/distance/neighbors/specializations/ivfpq_build_float_uint64_t.cu src/distance/neighbors/specializations/ivfpq_build_int8_t_uint64_t.cu src/distance/neighbors/specializations/ivfpq_build_uint8_t_uint64_t.cu diff --git a/cpp/include/raft/neighbors/ivf_flat.cuh b/cpp/include/raft/neighbors/ivf_flat.cuh index f18611b9f1..197e650671 100644 --- a/cpp/include/raft/neighbors/ivf_flat.cuh +++ b/cpp/include/raft/neighbors/ivf_flat.cuh @@ -102,8 +102,8 @@ auto build(raft::device_resources const& handle, * @tparam matrix_idx_t matrix indexing type * * @param[in] handle - * @param[in] params configure the index building * @param[in] dataset a device pointer to a row-major matrix [n_rows, dim] + * @param[in] params configure the index building * * @return the constructed ivf-flat index */ @@ -119,6 +119,52 @@ auto build(raft::device_resources const& handle, static_cast(dataset.extent(1))); } +/** + * @brief Build the index from the dataset for efficient search. + * + * NB: Currently, the following distance metrics are supported: + * - L2Expanded + * - L2Unexpanded + * - InnerProduct + * + * Usage example: + * @code{.cpp} + * using namespace raft::neighbors; + * // use default index parameters + * ivf_flat::index_params index_params; + * // create and fill the index from a [N, D] dataset + * ivf_flat::index *index_ptr; + * ivf_flat::build(handle, dataset, index_params, index_ptr); + * // use default search parameters + * ivf_flat::search_params search_params; + * // search K nearest neighbours for each of the N queries + * ivf_flat::search(handle, index, queries, out_inds, out_dists, search_params, k); + * @endcode + * + * @tparam value_t data element type + * @tparam idx_t type of the indices in the source dataset + * @tparam int_t precision / type of integral arguments + * @tparam matrix_idx_t matrix indexing type + * + * @param[in] handle + * @param[in] dataset a device pointer to a row-major matrix [n_rows, dim] + * @param[in] params configure the index building + * @param[out] idx pointer to ivf_flat::index + * + */ +template +void build(raft::device_resources const& handle, + raft::device_matrix_view dataset, + const index_params& params, + raft::neighbors::ivf_flat::index* idx) +{ + *idx = raft::spatial::knn::ivf_flat::detail::build(handle, + params, + dataset.data_handle(), + static_cast(dataset.extent(0)), + static_cast(dataset.extent(1))); +} + /** @} */ /** @@ -397,24 +443,21 @@ void search(raft::device_resources const& handle, * [n_queries, k] * @param[out] distances a device pointer to the distances to the selected neighbors [n_queries, k] * @param[in] params configure the search - * @param[in] k the number of neighbors to find for each query. */ -template +template void search(raft::device_resources const& handle, const index& index, raft::device_matrix_view queries, raft::device_matrix_view neighbors, raft::device_matrix_view distances, - const search_params& params, - int_t k) + const search_params& params) { 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) && neighbors.extent(1) == static_cast(k), - "Number of columns in output neighbors and distances matrices must equal k"); + RAFT_EXPECTS(neighbors.extent(1) == distances.extent(1) == neighbors.extent(1), + "Number of columns in output neighbors and distances matrices must be equal"); RAFT_EXPECTS(queries.extent(1) == index.dim(), "Number of query dimensions should equal number of dimensions in the index."); @@ -424,7 +467,7 @@ void search(raft::device_resources const& handle, index, queries.data_handle(), static_cast(queries.extent(0)), - static_cast(k), + static_cast(neighbors.extent(1)), neighbors.data_handle(), distances.data_handle(), nullptr); diff --git a/cpp/include/raft/neighbors/specializations.cuh b/cpp/include/raft/neighbors/specializations.cuh index 8652c21e4f..a6d3d9575e 100644 --- a/cpp/include/raft/neighbors/specializations.cuh +++ b/cpp/include/raft/neighbors/specializations.cuh @@ -21,6 +21,7 @@ #include #include +#include #include #include diff --git a/cpp/include/raft/neighbors/specializations/ivf_flat.cuh b/cpp/include/raft/neighbors/specializations/ivf_flat.cuh new file mode 100644 index 0000000000..ed4480c7fa --- /dev/null +++ b/cpp/include/raft/neighbors/specializations/ivf_flat.cuh @@ -0,0 +1,59 @@ +/* + * Copyright (c) 2023, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#pragma once + +#include + +namespace raft::neighbors::ivf_flat { + +#define RAFT_INST(T, IdxT) \ + extern template auto build(raft::device_resources const& handle, \ + raft::device_matrix_view dataset, \ + const index_params& params) \ + ->index; \ + \ + extern template auto extend( \ + raft::device_resources const& handle, \ + const index& orig_index, \ + raft::device_matrix_view new_vectors, \ + std::optional> new_indices) \ + ->index; \ + \ + extern template void build(raft::device_resources const& handle, \ + raft::device_matrix_view dataset, \ + const index_params& params, \ + index* idx); \ + \ + extern template void extend( \ + raft::device_resources const& handle, \ + index* idx, \ + raft::device_matrix_view new_vectors, \ + std::optional> new_indices); \ + \ + extern template void search(raft::device_resources const&, \ + const index&, \ + raft::device_matrix_view, \ + raft::device_matrix_view, \ + raft::device_matrix_view, \ + search_params const&); + +RAFT_INST(float, uint64_t); +RAFT_INST(int8_t, uint64_t); +RAFT_INST(uint8_t, uint64_t); + +#undef RAFT_INST +} // namespace raft::neighbors::ivf_flat diff --git a/cpp/include/raft_runtime/neighbors/ivf_flat.hpp b/cpp/include/raft_runtime/neighbors/ivf_flat.hpp index 989bfcd4c6..c4c751b182 100644 --- a/cpp/include/raft_runtime/neighbors/ivf_flat.hpp +++ b/cpp/include/raft_runtime/neighbors/ivf_flat.hpp @@ -57,8 +57,7 @@ RAFT_INST_BUILD_EXTEND(uint8_t, uint64_t) raft::device_matrix_view, \ raft::device_matrix_view, \ raft::device_matrix_view, \ - raft::neighbors::ivf_flat::search_params const&, \ - uint32_t); + raft::neighbors::ivf_flat::search_params const&); RAFT_INST_SEARCH(float, uint64_t); RAFT_INST_SEARCH(int8_t, uint64_t); diff --git a/cpp/src/distance/neighbors/ivf_flat_build.cu b/cpp/src/distance/neighbors/ivf_flat_build.cu index e31e81e070..02ab0fe653 100644 --- a/cpp/src/distance/neighbors/ivf_flat_build.cu +++ b/cpp/src/distance/neighbors/ivf_flat_build.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -#include +#include #include namespace raft::runtime::neighbors::ivf_flat { diff --git a/cpp/src/distance/neighbors/ivf_flat_search.cu b/cpp/src/distance/neighbors/ivf_flat_search.cu index 50441e25d8..dc060460be 100644 --- a/cpp/src/distance/neighbors/ivf_flat_search.cu +++ b/cpp/src/distance/neighbors/ivf_flat_search.cu @@ -14,8 +14,7 @@ * limitations under the License. */ -#include -// #include +#include #include namespace raft::runtime::neighbors::ivf_flat { @@ -26,11 +25,10 @@ namespace raft::runtime::neighbors::ivf_flat { raft::device_matrix_view queries, \ raft::device_matrix_view neighbors, \ raft::device_matrix_view distances, \ - raft::neighbors::ivf_flat::search_params const& params, \ - uint32_t k) \ + raft::neighbors::ivf_flat::search_params const& params) \ { \ - raft::neighbors::ivf_flat::search( \ - handle, index, queries, neighbors, distances, params, k); \ + raft::neighbors::ivf_flat::search( \ + handle, index, queries, neighbors, distances, params); \ } RAFT_INST_SEARCH(float, uint64_t); diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_build_float_uint64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_build_float_uint64_t.cu new file mode 100644 index 0000000000..3cb561fdc0 --- /dev/null +++ b/cpp/src/distance/neighbors/specializations/ivfflat_build_float_uint64_t.cu @@ -0,0 +1,36 @@ +/* + * Copyright (c) 2023, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +namespace raft::neighbors::ivf_flat { + +#define RAFT_MAKE_INSTANCE(T, IdxT) \ + template auto build(raft::device_resources const& handle, \ + raft::device_matrix_view dataset, \ + const index_params& params) \ + ->index; \ + \ + template void build(raft::device_resources const& handle, \ + raft::device_matrix_view dataset, \ + const index_params& params, \ + index* idx); + +RAFT_MAKE_INSTANCE(float, uint64_t); + +#undef RAFT_MAKE_INSTANCE + +} // namespace raft::neighbors::ivf_flat diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_build_int8_t_uint64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_build_int8_t_uint64_t.cu new file mode 100644 index 0000000000..babca2a32c --- /dev/null +++ b/cpp/src/distance/neighbors/specializations/ivfflat_build_int8_t_uint64_t.cu @@ -0,0 +1,36 @@ +/* + * Copyright (c) 2023, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +namespace raft::neighbors::ivf_flat { + +#define RAFT_MAKE_INSTANCE(T, IdxT) \ + template auto build(raft::device_resources const& handle, \ + raft::device_matrix_view dataset, \ + const index_params& params) \ + ->index; \ + \ + template void build(raft::device_resources const& handle, \ + raft::device_matrix_view dataset, \ + const index_params& params, \ + index* idx); + +RAFT_MAKE_INSTANCE(int8_t, uint64_t); + +#undef RAFT_MAKE_INSTANCE + +} // namespace raft::neighbors::ivf_flat diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_build_uint8_t_uint64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_build_uint8_t_uint64_t.cu new file mode 100644 index 0000000000..c18b3dd696 --- /dev/null +++ b/cpp/src/distance/neighbors/specializations/ivfflat_build_uint8_t_uint64_t.cu @@ -0,0 +1,36 @@ +/* + * Copyright (c) 2023, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +namespace raft::neighbors::ivf_flat { + +#define RAFT_MAKE_INSTANCE(T, IdxT) \ + template auto build(raft::device_resources const& handle, \ + raft::device_matrix_view dataset, \ + const index_params& params) \ + ->index; \ + \ + template void build(raft::device_resources const& handle, \ + raft::device_matrix_view dataset, \ + const index_params& params, \ + index* idx); + +RAFT_MAKE_INSTANCE(uint8_t, uint64_t); + +#undef RAFT_MAKE_INSTANCE + +} // namespace raft::neighbors::ivf_flat diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_extend_float_uint64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_extend_float_uint64_t.cu new file mode 100644 index 0000000000..614a7a6a10 --- /dev/null +++ b/cpp/src/distance/neighbors/specializations/ivfflat_extend_float_uint64_t.cu @@ -0,0 +1,37 @@ +/* + * Copyright (c) 2023, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +namespace raft::neighbors::ivf_flat { + +#define RAFT_MAKE_INSTANCE(T, IdxT) \ + template auto extend(raft::device_resources const& handle, \ + const raft::neighbors::ivf_flat::index& orig_index, \ + raft::device_matrix_view new_vectors, \ + std::optional> new_indices) \ + ->raft::neighbors::ivf_flat::index; \ + \ + template void extend(raft::device_resources const& handle, \ + raft::neighbors::ivf_flat::index* idx, \ + raft::device_matrix_view new_vectors, \ + std::optional> new_indices); + +RAFT_MAKE_INSTANCE(float, uint64_t); + +#undef RAFT_MAKE_INSTANCE + +} // namespace raft::neighbors::ivf_flat diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_extend_int8_t_uint64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_extend_int8_t_uint64_t.cu new file mode 100644 index 0000000000..fe0f63a406 --- /dev/null +++ b/cpp/src/distance/neighbors/specializations/ivfflat_extend_int8_t_uint64_t.cu @@ -0,0 +1,37 @@ +/* + * Copyright (c) 2023, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +namespace raft::neighbors::ivf_flat { + +#define RAFT_MAKE_INSTANCE(T, IdxT) \ + template auto extend(raft::device_resources const& handle, \ + const raft::neighbors::ivf_flat::index& orig_index, \ + raft::device_matrix_view new_vectors, \ + std::optional> new_indices) \ + ->raft::neighbors::ivf_flat::index; \ + \ + template void extend(raft::device_resources const& handle, \ + raft::neighbors::ivf_flat::index* idx, \ + raft::device_matrix_view new_vectors, \ + std::optional> new_indices); + +RAFT_MAKE_INSTANCE(int8_t, uint64_t); + +#undef RAFT_MAKE_INSTANCE + +} // namespace raft::neighbors::ivf_flat diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_extend_uint8_t_uint64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_extend_uint8_t_uint64_t.cu new file mode 100644 index 0000000000..187d47f842 --- /dev/null +++ b/cpp/src/distance/neighbors/specializations/ivfflat_extend_uint8_t_uint64_t.cu @@ -0,0 +1,36 @@ +/* + * Copyright (c) 2023, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +namespace raft::neighbors::ivf_flat { + +#define RAFT_MAKE_INSTANCE(T, IdxT) \ + template auto extend(raft::device_resources const& handle, \ + const raft::neighbors::ivf_flat::index& orig_index, \ + raft::device_matrix_view new_vectors, \ + std::optional> new_indices) \ + ->raft::neighbors::ivf_flat::index; \ + \ + template void extend(raft::device_resources const& handle, \ + raft::neighbors::ivf_flat::index* idx, \ + raft::device_matrix_view new_vectors, \ + std::optional> new_indices); +RAFT_MAKE_INSTANCE(uint8_t, uint64_t); + +#undef RAFT_MAKE_INSTANCE + +} // namespace raft::neighbors::ivf_flat diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_search_float_uint64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_search_float_uint64_t.cu new file mode 100644 index 0000000000..fcc3b120fc --- /dev/null +++ b/cpp/src/distance/neighbors/specializations/ivfflat_search_float_uint64_t.cu @@ -0,0 +1,33 @@ +/* + * Copyright (c) 2023, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +namespace raft::neighbors::ivf_flat { + +#define RAFT_MAKE_INSTANCE(T, IdxT) \ + template void search(raft::device_resources const&, \ + const raft::neighbors::ivf_flat::index&, \ + raft::device_matrix_view, \ + raft::device_matrix_view, \ + raft::device_matrix_view, \ + raft::neighbors::ivf_flat::search_params const&); + +RAFT_MAKE_INSTANCE(float, uint64_t); + +#undef RAFT_MAKE_INSTANCE + +} // namespace raft::neighbors::ivf_flat diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_search_int8_t_uint64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_search_int8_t_uint64_t.cu new file mode 100644 index 0000000000..12220ebb5a --- /dev/null +++ b/cpp/src/distance/neighbors/specializations/ivfflat_search_int8_t_uint64_t.cu @@ -0,0 +1,33 @@ +/* + * Copyright (c) 2023, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +namespace raft::neighbors::ivf_flat { + +#define RAFT_MAKE_INSTANCE(T, IdxT) \ + template void search(raft::device_resources const&, \ + const raft::neighbors::ivf_flat::index&, \ + raft::device_matrix_view, \ + raft::device_matrix_view, \ + raft::device_matrix_view, \ + raft::neighbors::ivf_flat::search_params const&); + +RAFT_MAKE_INSTANCE(int8_t, uint64_t); + +#undef RAFT_MAKE_INSTANCE + +} // namespace raft::neighbors::ivf_flat diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_search_uint8_t_uint64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_search_uint8_t_uint64_t.cu new file mode 100644 index 0000000000..89859137f6 --- /dev/null +++ b/cpp/src/distance/neighbors/specializations/ivfflat_search_uint8_t_uint64_t.cu @@ -0,0 +1,33 @@ +/* + * Copyright (c) 2023, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include + +namespace raft::neighbors::ivf_flat { + +#define RAFT_MAKE_INSTANCE(T, IdxT) \ + template void search(raft::device_resources const&, \ + const raft::neighbors::ivf_flat::index&, \ + raft::device_matrix_view, \ + raft::device_matrix_view, \ + raft::device_matrix_view, \ + raft::neighbors::ivf_flat::search_params const&); + +RAFT_MAKE_INSTANCE(uint8_t, uint64_t); + +#undef RAFT_MAKE_INSTANCE + +} // namespace raft::neighbors::ivf_flat diff --git a/cpp/test/neighbors/ann_ivf_flat.cuh b/cpp/test/neighbors/ann_ivf_flat.cuh index 26b8301cb1..8d6b184c19 100644 --- a/cpp/test/neighbors/ann_ivf_flat.cuh +++ b/cpp/test/neighbors/ann_ivf_flat.cuh @@ -202,8 +202,7 @@ class AnnIVFFlatTest : public ::testing::TestWithParam> { search_queries_view, indices_out_view, dists_out_view, - search_params, - ps.k); + search_params); update_host(distances_ivfflat.data(), distances_ivfflat_dev.data(), queries_size, stream_); update_host(indices_ivfflat.data(), indices_ivfflat_dev.data(), queries_size, stream_); diff --git a/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/c_ivf_flat.pxd b/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/c_ivf_flat.pxd index 577612593a..4e62d42cd3 100644 --- a/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/c_ivf_flat.pxd +++ b/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/c_ivf_flat.pxd @@ -122,8 +122,7 @@ cdef extern from "raft_runtime/neighbors/ivf_flat.hpp" \ device_matrix_view[float, uint64_t, row_major] queries, device_matrix_view[uint64_t, uint64_t, row_major] neighbors, device_matrix_view[float, uint64_t, row_major] distances, - const search_params& params, - uint32_t k) except + + const search_params& params) except + cdef void search( const device_resources& handle, @@ -131,8 +130,7 @@ cdef extern from "raft_runtime/neighbors/ivf_flat.hpp" \ device_matrix_view[int8_t, uint64_t, row_major] queries, device_matrix_view[uint64_t, uint64_t, row_major] neighbors, device_matrix_view[float, uint64_t, row_major] distances, - const search_params& params, - uint32_t k) except + + const search_params& params) except + cdef void search( const device_resources& handle, @@ -140,8 +138,7 @@ cdef extern from "raft_runtime/neighbors/ivf_flat.hpp" \ device_matrix_view[uint8_t, uint64_t, row_major] queries, device_matrix_view[uint64_t, uint64_t, row_major] neighbors, device_matrix_view[float, uint64_t, row_major] distances, - const search_params& params, - uint32_t k) except + + const search_params& params) except + # cdef void save(const device_resources& handle, # const string& filename, diff --git a/python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx b/python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx index 16b259c7c1..f777393c20 100644 --- a/python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx +++ b/python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx @@ -598,8 +598,7 @@ def search(SearchParams search_params, get_dmv_float(queries, check_shape=True), get_dmv_uint64(neighbors, check_shape=True), get_dmv_float(distances, check_shape=True), - params, - k) + params) elif queries_dt == np.byte: with cuda_interruptible(): c_ivf_flat.search(deref(handle_), @@ -607,8 +606,7 @@ def search(SearchParams search_params, get_dmv_int8(queries, check_shape=True), get_dmv_uint64(neighbors, check_shape=True), get_dmv_float(distances, check_shape=True), - params, - k) + params) elif queries_dt == np.ubyte: with cuda_interruptible(): c_ivf_flat.search(deref(handle_), @@ -616,8 +614,7 @@ def search(SearchParams search_params, get_dmv_uint8(queries, check_shape=True), get_dmv_uint64(neighbors, check_shape=True), get_dmv_float(distances, check_shape=True), - params, - k) + params) else: raise ValueError("query dtype %s not supported" % queries_dt) From 6925e842ca4323ccb43dca35bd1ff57bd4d90668 Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 9 Mar 2023 10:00:00 -0800 Subject: [PATCH 12/20] consolidate python wrapper, write pytests --- .gitignore | 1 + cpp/include/raft/neighbors/ivf_flat.cuh | 2 +- .../pylibraft/neighbors/CMakeLists.txt | 2 +- .../pylibraft/pylibraft/neighbors/__init__.py | 4 +- .../pylibraft/pylibraft/neighbors/common.pxd | 24 + .../pylibraft/pylibraft/neighbors/common.pyx | 61 +++ .../pylibraft/neighbors/ivf_flat/ivf_flat.pyx | 237 ++++++--- .../pylibraft/neighbors/ivf_pq/ivf_pq.pyx | 43 +- .../pylibraft/pylibraft/test/test_ivf_flat.py | 463 ++++++++++++++++++ 9 files changed, 723 insertions(+), 114 deletions(-) create mode 100644 python/pylibraft/pylibraft/neighbors/common.pxd create mode 100644 python/pylibraft/pylibraft/neighbors/common.pyx create mode 100644 python/pylibraft/pylibraft/test/test_ivf_flat.py diff --git a/.gitignore b/.gitignore index 80709dbb96..c2528d2cd0 100644 --- a/.gitignore +++ b/.gitignore @@ -24,6 +24,7 @@ log .DS_Store dask-worker-space/ *.egg-info/ +*.bin ## scikit-build _skbuild diff --git a/cpp/include/raft/neighbors/ivf_flat.cuh b/cpp/include/raft/neighbors/ivf_flat.cuh index 197e650671..109aaee933 100644 --- a/cpp/include/raft/neighbors/ivf_flat.cuh +++ b/cpp/include/raft/neighbors/ivf_flat.cuh @@ -456,7 +456,7 @@ void search(raft::device_resources const& handle, 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) == neighbors.extent(1), + RAFT_EXPECTS(neighbors.extent(1) == distances.extent(1), "Number of columns in output neighbors and distances matrices must be equal"); RAFT_EXPECTS(queries.extent(1) == index.dim(), diff --git a/python/pylibraft/pylibraft/neighbors/CMakeLists.txt b/python/pylibraft/pylibraft/neighbors/CMakeLists.txt index 6fef9288d3..572ea47f4e 100644 --- a/python/pylibraft/pylibraft/neighbors/CMakeLists.txt +++ b/python/pylibraft/pylibraft/neighbors/CMakeLists.txt @@ -13,7 +13,7 @@ # ============================================================================= # Set the list of Cython files to build -set(cython_sources refine.pyx) +set(cython_sources common.pyx refine.pyx) set(linked_libraries raft::raft raft::distance) # Build all of the Cython targets diff --git a/python/pylibraft/pylibraft/neighbors/__init__.py b/python/pylibraft/pylibraft/neighbors/__init__.py index dd8cdd8445..f7510ba2db 100644 --- a/python/pylibraft/pylibraft/neighbors/__init__.py +++ b/python/pylibraft/pylibraft/neighbors/__init__.py @@ -1,4 +1,4 @@ -# Copyright (c) 2022, NVIDIA CORPORATION. +# Copyright (c) 2022-2023, NVIDIA CORPORATION. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -14,4 +14,4 @@ # from .refine import refine -__all__ = ["refine"] +__all__ = ["common", "refine"] diff --git a/python/pylibraft/pylibraft/neighbors/common.pxd b/python/pylibraft/pylibraft/neighbors/common.pxd new file mode 100644 index 0000000000..b11ef3176e --- /dev/null +++ b/python/pylibraft/pylibraft/neighbors/common.pxd @@ -0,0 +1,24 @@ +# +# Copyright (c) 2023, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# cython: profile=False +# distutils: language = c++ +# cython: embedsignature = True +# cython: language_level = 3 + +from pylibraft.distance.distance_type cimport DistanceType + + +cdef _get_metric_string(DistanceType metric) diff --git a/python/pylibraft/pylibraft/neighbors/common.pyx b/python/pylibraft/pylibraft/neighbors/common.pyx new file mode 100644 index 0000000000..1de782c185 --- /dev/null +++ b/python/pylibraft/pylibraft/neighbors/common.pyx @@ -0,0 +1,61 @@ +# +# Copyright (c) 2023, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# +# cython: profile=False +# distutils: language = c++ +# cython: embedsignature = True +# cython: language_level = 3 + +import warnings + +from pylibraft.distance.distance_type cimport DistanceType + + +def _get_metric(metric): + SUPPORTED_DISTANCES = { + "sqeuclidean": DistanceType.L2Expanded, + "euclidean": DistanceType.L2SqrtExpanded, + "inner_product": DistanceType.InnerProduct + } + if metric not in SUPPORTED_DISTANCES: + if metric == "l2_expanded": + warnings.warn("Using l2_expanded as a metric name is deprecated," + " use sqeuclidean instead", FutureWarning) + return DistanceType.L2Expanded + + raise ValueError("metric %s is not supported" % metric) + return SUPPORTED_DISTANCES[metric] + + +cdef _get_metric_string(DistanceType metric): + return {DistanceType.L2Expanded : "sqeuclidean", + DistanceType.InnerProduct: "inner_product", + DistanceType.L2SqrtExpanded: "euclidean"}[metric] + + +def _check_input_array(cai, exp_dt, exp_rows=None, exp_cols=None): + if cai.dtype not in exp_dt: + raise TypeError("dtype %s not supported" % cai["typestr"]) + + if not cai.c_contiguous: + raise ValueError("Row major input is expected") + + if exp_cols is not None and cai.shape[1] != exp_cols: + raise ValueError("Incorrect number of columns, expected {} got {}" + .format(exp_cols, cai.shape[1])) + + if exp_rows is not None and cai.shape[0] != exp_rows: + raise ValueError("Incorrect number of rows, expected {} , got {}" + .format(exp_rows, cai.shape[0])) diff --git a/python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx b/python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx index f777393c20..ea0618fa4a 100644 --- a/python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx +++ b/python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx @@ -40,7 +40,6 @@ from pylibraft.common import ( DeviceResources, ai_wrapper, auto_convert_output, - cai_wrapper, device_ndarray, ) from pylibraft.common.cai_wrapper import cai_wrapper @@ -64,21 +63,18 @@ from rmm._lib.memory_resource cimport ( device_memory_resource, ) +cimport pylibraft.neighbors.ivf_flat.cpp.c_ivf_flat as c_ivf_flat from pylibraft.common.cpp.optional cimport optional -from pylibraft.neighbors.ivf_pq import ( - _check_input_array, - _get_metric, - _get_metric_string, -) +from pylibraft.neighbors.common import _check_input_array, _get_metric -cimport pylibraft.neighbors.ivf_flat.cpp.c_ivf_flat as c_ivf_flat from pylibraft.common.mdspan cimport ( get_dmv_float, get_dmv_int8, get_dmv_uint8, get_dmv_uint64, ) +from pylibraft.neighbors.common cimport _get_metric_string from pylibraft.neighbors.ivf_flat.cpp.c_ivf_flat cimport ( index_params, search_params, @@ -168,49 +164,128 @@ cdef class IndexParams: cdef class Index: - # We store a pointer to the index because it dose not have a trivial - # constructor. - cdef c_ivf_flat.index[float, uint64_t] * index_float - cdef c_ivf_flat.index[int8_t, uint64_t] * index_int8 - cdef c_ivf_flat.index[uint8_t, uint64_t] * index_uint8 - cdef readonly bool trained + cdef str active_index_type - def __cinit__(self, handle=None): + def __cinit__(self): self.trained = False + self.active_index_type = None + + +cdef class IndexFloat(Index): + cdef c_ivf_flat.index[float, uint64_t] * index + + def __cinit__(self, handle=None): if handle is None: handle = DeviceResources() cdef device_resources* handle_ = \ handle.getHandle() + # this is to keep track of which index type is being used # We create a placeholder object. The actual parameter values do # not matter, it will be replaced with a built index object later. - self.index_float = new c_ivf_flat.index[float, uint64_t]( + self.index = new c_ivf_flat.index[float, uint64_t]( deref(handle_), _get_metric("sqeuclidean"), 1, False, 8) - self.index_int8 = new c_ivf_flat.index[int8_t, uint64_t]( + def __repr__(self): + m_str = "metric=" + _get_metric_string(self.index.metric()) + attr_str = [ + attr + "=" + str(getattr(self, attr)) + for attr in ["size", "dim", "n_lists", "adaptive_centers"] + ] + attr_str = [m_str] + attr_str + return "Index(type=IVF-FLAT, " + (", ".join(attr_str)) + ")" + + @property + def dim(self): + return self.index[0].dim() + + @property + def size(self): + return self.index[0].size() + + @property + def metric(self): + return self.index[0].metric() + + @property + def n_lists(self): + return self.index[0].n_lists() + + @property + def adaptive_centers(self): + return self.index[0].adaptive_centers() + + +cdef class IndexInt8(Index): + cdef c_ivf_flat.index[int8_t, uint64_t] * index + + def __cinit__(self, handle=None): + if handle is None: + handle = DeviceResources() + cdef device_resources* handle_ = \ + handle.getHandle() + + # this is to keep track of which index type is being used + # We create a placeholder object. The actual parameter values do + # not matter, it will be replaced with a built index object later. + self.index = new c_ivf_flat.index[int8_t, uint64_t]( deref(handle_), _get_metric("sqeuclidean"), 1, False, 8) - self.index_uint8 = new c_ivf_flat.index[uint8_t, uint64_t]( + def __repr__(self): + m_str = "metric=" + _get_metric_string(self.index.metric()) + attr_str = [ + attr + "=" + str(getattr(self, attr)) + for attr in ["size", "dim", "n_lists", "adaptive_centers"] + ] + attr_str = [m_str] + attr_str + return "Index(type=IVF-FLAT, " + (", ".join(attr_str)) + ")" + + @property + def dim(self): + return self.index[0].dim() + + @property + def size(self): + return self.index[0].size() + + @property + def metric(self): + return self.index[0].metric() + + @property + def n_lists(self): + return self.index[0].n_lists() + + @property + def adaptive_centers(self): + return self.index[0].adaptive_centers() + + +cdef class IndexUint8(Index): + cdef c_ivf_flat.index[uint8_t, uint64_t] * index + + def __cinit__(self, handle=None): + if handle is None: + handle = DeviceResources() + cdef device_resources* handle_ = \ + handle.getHandle() + + # this is to keep track of which index type is being used + # We create a placeholder object. The actual parameter values do + # not matter, it will be replaced with a built index object later. + self.index = new c_ivf_flat.index[uint8_t, uint64_t]( deref(handle_), _get_metric("sqeuclidean"), 1, False, 8) - def __dealloc__(self): - if self.index_float is not NULL: - del self.index_float - if self.index_int8 is not NULL: - del self.index_int8 - if self.index_uint8 is not NULL: - del self.index_uint8 - def __repr__(self): m_str = "metric=" + _get_metric_string(self.index.metric()) attr_str = [ @@ -306,32 +381,43 @@ def build(IndexParams index_params, dataset, handle=None): cdef device_resources* handle_ = \ handle.getHandle() - idx = Index() + cdef IndexFloat idx_float + cdef IndexInt8 idx_int8 + cdef IndexUint8 idx_uint8 if dataset_dt == np.float32: + idx_float = IndexFloat(handle) + idx_float.active_index_type = "float32" with cuda_interruptible(): c_ivf_flat.build(deref(handle_), - get_dmv_float(dataset, check_shape=True), + get_dmv_float(dataset_cai, check_shape=True), index_params.params, - idx.index_float) + idx_float.index) + idx_float.trained = True + return idx_float elif dataset_dt == np.byte: + idx_int8 = IndexInt8(handle) + idx_int8.active_index_type = "byte" with cuda_interruptible(): c_ivf_flat.build(deref(handle_), - get_dmv_int8(dataset, check_shape=True), + get_dmv_int8(dataset_cai, check_shape=True), index_params.params, - idx.index_int8) + idx_int8.index) + idx_int8.trained = True + return idx_int8 elif dataset_dt == np.ubyte: + idx_uint8 = IndexUint8(handle) + idx_uint8.active_index_type = "ubyte" with cuda_interruptible(): c_ivf_flat.build(deref(handle_), - get_dmv_uint8(dataset, check_shape=True), + get_dmv_uint8(dataset_cai, check_shape=True), index_params.params, - idx.index_uint8) + idx_uint8.index) + idx_uint8.trained = True + return idx_uint8 else: raise TypeError("dtype %s not supported" % dataset_dt) - idx.trained = True - return idx - @auto_sync_handle @auto_convert_output @@ -404,8 +490,7 @@ def extend(Index index, new_vectors, new_indices, handle=None): cdef uint64_t n_rows = vecs_cai.shape[0] cdef uint32_t dim = vecs_cai.shape[1] - _check_input_array(vecs_cai, [np.dtype('float32'), np.dtype('byte'), - np.dtype('ubyte')], + _check_input_array(vecs_cai, [np.dtype(index.active_index_type)], exp_cols=index.dim) idx_cai = cai_wrapper(new_indices) @@ -413,40 +498,45 @@ def extend(Index index, new_vectors, new_indices, handle=None): if len(idx_cai.shape)!=1: raise ValueError("Indices array is expected to be 1D") - cdef optional[device_vector_view[uint64_t, uint64_t]] new_indices_float - cdef optional[device_vector_view[uint64_t, uint64_t]] new_indices_int8 - cdef optional[device_vector_view[uint64_t, uint64_t]] new_indices_uint8 + cdef optional[device_vector_view[uint64_t, uint64_t]] new_indices_opt + + cdef IndexFloat idx_float + cdef IndexInt8 idx_int8 + cdef IndexUint8 idx_uint8 if vecs_dt == np.float32: - if index.index_float[0].size() > 0: - new_indices_float = make_device_vector_view( + idx_float = index + if idx_float.index.size() > 0: + new_indices_opt = make_device_vector_view( idx_cai.data, idx_cai.shape[0]) with cuda_interruptible(): c_ivf_flat.extend(deref(handle_), - index.index_float, - get_dmv_float(new_vectors, check_shape=True), - new_indices_float) + idx_float.index, + get_dmv_float(vecs_cai, check_shape=True), + new_indices_opt) elif vecs_dt == np.int8: - if index.index_int8[0].size() > 0: - new_indices_int8 = make_device_vector_view( + idx_int8 = index + if idx_int8.index[0].size() > 0: + new_indices_opt = make_device_vector_view( idx_cai.data, idx_cai.shape[0]) with cuda_interruptible(): c_ivf_flat.extend(deref(handle_), - index.index_int8, - get_dmv_int8(new_vectors, check_shape=True), - new_indices_int8) + idx_int8.index, + get_dmv_int8(vecs_cai, check_shape=True), + new_indices_opt) elif vecs_dt == np.uint8: - if index.index_uint8[0].size() > 0: - new_indices_uint8 = make_device_vector_view( + idx_uint8 = index + if idx_uint8.index[0].size() > 0: + new_indices_opt = make_device_vector_view( idx_cai.data, idx_cai.shape[0]) with cuda_interruptible(): c_ivf_flat.extend(deref(handle_), - index.index_uint8, - get_dmv_uint8(new_vectors, check_shape=True), - new_indices_uint8) + idx_uint8.index, + get_dmv_uint8(vecs_cai, check_shape=True), + new_indices_opt) else: raise TypeError("query dtype %s not supported" % vecs_dt) @@ -456,9 +546,7 @@ def extend(Index index, new_vectors, new_indices, handle=None): cdef class SearchParams: cdef c_ivf_flat.search_params params - def __init__(self, *, n_probes=20, - lut_dtype=np.float32, - internal_distance_dtype=np.float32): + def __init__(self, *, n_probes=20): """ IVF-FLAT search parameters @@ -571,8 +659,7 @@ def search(SearchParams search_params, queries_dt = queries_cai.dtype cdef uint32_t n_queries = queries_cai.shape[0] - _check_input_array(queries_cai, [np.dtype('float32'), np.dtype('byte'), - np.dtype('ubyte')], + _check_input_array(queries_cai, [np.dtype(index.active_index_type)], exp_cols=index.dim) if neighbors is None: @@ -590,30 +677,36 @@ def search(SearchParams search_params, exp_rows=n_queries, exp_cols=k) cdef c_ivf_flat.search_params params = search_params.params + cdef IndexFloat idx_float + cdef IndexInt8 idx_int8 + cdef IndexUint8 idx_uint8 if queries_dt == np.float32: + idx_float = index with cuda_interruptible(): c_ivf_flat.search(deref(handle_), - deref(index.index_float), - get_dmv_float(queries, check_shape=True), - get_dmv_uint64(neighbors, check_shape=True), - get_dmv_float(distances, check_shape=True), + deref(idx_float.index), + get_dmv_float(queries_cai, check_shape=True), + get_dmv_uint64(neighbors_cai, check_shape=True), + get_dmv_float(distances_cai, check_shape=True), params) elif queries_dt == np.byte: + idx_int8 = index with cuda_interruptible(): c_ivf_flat.search(deref(handle_), - deref(index.index_int8), - get_dmv_int8(queries, check_shape=True), - get_dmv_uint64(neighbors, check_shape=True), - get_dmv_float(distances, check_shape=True), + deref(idx_int8.index), + get_dmv_int8(queries_cai, check_shape=True), + get_dmv_uint64(neighbors_cai, check_shape=True), + get_dmv_float(distances_cai, check_shape=True), params) elif queries_dt == np.ubyte: + idx_uint8 = index with cuda_interruptible(): c_ivf_flat.search(deref(handle_), - deref(index.index_uint8), - get_dmv_uint8(queries, check_shape=True), - get_dmv_uint64(neighbors, check_shape=True), - get_dmv_float(distances, check_shape=True), + deref(idx_uint8.index), + get_dmv_uint8(queries_cai, check_shape=True), + get_dmv_uint64(neighbors_cai, check_shape=True), + get_dmv_float(distances_cai, check_shape=True), params) else: raise ValueError("query dtype %s not supported" % queries_dt) diff --git a/python/pylibraft/pylibraft/neighbors/ivf_pq/ivf_pq.pyx b/python/pylibraft/pylibraft/neighbors/ivf_pq/ivf_pq.pyx index 47d8e94e5f..39d1fe0514 100644 --- a/python/pylibraft/pylibraft/neighbors/ivf_pq/ivf_pq.pyx +++ b/python/pylibraft/pylibraft/neighbors/ivf_pq/ivf_pq.pyx @@ -49,7 +49,11 @@ from rmm._lib.memory_resource cimport ( device_memory_resource, ) +cimport pylibraft.neighbors.ivf_flat.cpp.c_ivf_flat as c_ivf_flat cimport pylibraft.neighbors.ivf_pq.cpp.c_ivf_pq as c_ivf_pq + +from pylibraft.neighbors.common import _check_input_array, _get_metric + from pylibraft.common.cpp.mdspan cimport device_matrix_view from pylibraft.common.mdspan cimport ( get_dmv_float, @@ -57,34 +61,13 @@ from pylibraft.common.mdspan cimport ( get_dmv_uint8, get_dmv_uint64, ) +from pylibraft.neighbors.common cimport _get_metric_string from pylibraft.neighbors.ivf_pq.cpp.c_ivf_pq cimport ( index_params, search_params, ) -def _get_metric(metric): - SUPPORTED_DISTANCES = { - "sqeuclidean": DistanceType.L2Expanded, - "euclidean": DistanceType.L2SqrtExpanded, - "inner_product": DistanceType.InnerProduct - } - if metric not in SUPPORTED_DISTANCES: - if metric == "l2_expanded": - warnings.warn("Using l2_expanded as a metric name is deprecated," - " use sqeuclidean instead", FutureWarning) - return DistanceType.L2Expanded - - raise ValueError("metric %s is not supported" % metric) - return SUPPORTED_DISTANCES[metric] - - -cdef _get_metric_string(DistanceType metric): - return {DistanceType.L2Expanded : "sqeuclidean", - DistanceType.InnerProduct: "inner_product", - DistanceType.L2SqrtExpanded: "euclidean"}[metric] - - cdef _get_codebook_string(c_ivf_pq.codebook_gen codebook): return {c_ivf_pq.codebook_gen.PER_SUBSPACE: "subspace", c_ivf_pq.codebook_gen.PER_CLUSTER: "cluster"}[codebook] @@ -104,22 +87,6 @@ cdef _get_dtype_string(dtype): c_ivf_pq.cudaDataType_t.CUDA_R_8U: np.uint8}[dtype]) -def _check_input_array(cai, exp_dt, exp_rows=None, exp_cols=None): - if cai.dtype not in exp_dt: - raise TypeError("dtype %s not supported" % cai["typestr"]) - - if not cai.c_contiguous: - raise ValueError("Row major input is expected") - - if exp_cols is not None and cai.shape[1] != exp_cols: - raise ValueError("Incorrect number of columns, expected {} got {}" - .format(exp_cols, cai.shape[1])) - - if exp_rows is not None and cai.shape[0] != exp_rows: - raise ValueError("Incorrect number of rows, expected {} , got {}" - .format(exp_rows, cai.shape[0])) - - cdef class IndexParams: cdef c_ivf_pq.index_params params diff --git a/python/pylibraft/pylibraft/test/test_ivf_flat.py b/python/pylibraft/pylibraft/test/test_ivf_flat.py new file mode 100644 index 0000000000..ce6d3c6a4b --- /dev/null +++ b/python/pylibraft/pylibraft/test/test_ivf_flat.py @@ -0,0 +1,463 @@ +# Copyright (c) 2023, NVIDIA CORPORATION. +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# 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.metrics import pairwise_distances +from sklearn.neighbors import NearestNeighbors +from sklearn.preprocessing import normalize + +from pylibraft.common import device_ndarray +from pylibraft.neighbors import ivf_flat + + +def generate_data(shape, dtype): + if dtype == np.byte: + x = np.random.randint(-127, 128, size=shape, dtype=np.byte) + elif dtype == np.ubyte: + x = np.random.randint(0, 255, size=shape, dtype=np.ubyte) + else: + x = np.random.random_sample(shape).astype(dtype) + + return x + + +def calc_recall(ann_idx, true_nn_idx): + assert ann_idx.shape == true_nn_idx.shape + n = 0 + for i in range(ann_idx.shape[0]): + n += np.intersect1d(ann_idx[i, :], true_nn_idx[i, :]).size + recall = n / ann_idx.size + return recall + + +def check_distances(dataset, queries, metric, out_idx, out_dist, eps=None): + """ + Calculate the real distance between queries and dataset[out_idx], + and compare it to out_dist. + """ + if eps is None: + # Quantization leads to errors in the distance calculation. + # The aim of this test is not to test precision, but to catch obvious + # errors. + eps = 0.1 + + dist = np.empty(out_dist.shape, out_dist.dtype) + for i in range(queries.shape[0]): + X = queries[np.newaxis, i, :] + Y = dataset[out_idx[i, :], :] + if metric == "sqeuclidean": + dist[i, :] = pairwise_distances(X, Y, "sqeuclidean") + elif metric == "euclidean": + dist[i, :] = pairwise_distances(X, Y, "euclidean") + elif metric == "inner_product": + dist[i, :] = np.matmul(X, Y.T) + else: + raise ValueError("Invalid metric") + + dist_eps = abs(dist) + dist_eps[dist < 1e-3] = 1e-3 + diff = abs(out_dist - dist) / dist_eps + + assert np.mean(diff) < eps + + +def run_ivf_flat_build_search_test( + n_rows, + n_cols, + n_queries, + k, + n_lists, + metric, + dtype, + add_data_on_build=True, + n_probes=100, + kmeans_trainset_fraction=1, + kmeans_n_iters=20, + compare=True, + inplace=True, + array_type="device", +): + dataset = generate_data((n_rows, n_cols), dtype) + if metric == "inner_product": + dataset = normalize(dataset, norm="l2", axis=1) + dataset_device = device_ndarray(dataset) + + build_params = ivf_flat.IndexParams( + n_lists=n_lists, + metric=metric, + kmeans_n_iters=kmeans_n_iters, + kmeans_trainset_fraction=kmeans_trainset_fraction, + add_data_on_build=add_data_on_build, + ) + + if array_type == "device": + index = ivf_flat.build(build_params, dataset_device) + else: + index = ivf_flat.build(build_params, dataset) + + assert index.trained + + assert index.metric == build_params.metric + assert index.n_lists == build_params.n_lists + + if not add_data_on_build: + dataset_1 = dataset[: n_rows // 2, :] + dataset_2 = dataset[n_rows // 2 :, :] + indices_1 = np.arange(n_rows // 2, dtype=np.uint64) + indices_2 = np.arange(n_rows // 2, n_rows, dtype=np.uint64) + if array_type == "device": + dataset_1_device = device_ndarray(dataset_1) + dataset_2_device = device_ndarray(dataset_2) + indices_1_device = device_ndarray(indices_1) + indices_2_device = device_ndarray(indices_2) + index = ivf_flat.extend(index, dataset_1_device, indices_1_device) + index = ivf_flat.extend(index, dataset_2_device, indices_2_device) + else: + index = ivf_flat.extend(index, dataset_1, indices_1) + index = ivf_flat.extend(index, dataset_2, indices_2) + + assert index.size >= n_rows + + queries = generate_data((n_queries, n_cols), dtype) + out_idx = np.zeros((n_queries, k), dtype=np.uint64) + out_dist = np.zeros((n_queries, k), dtype=np.float32) + + queries_device = device_ndarray(queries) + out_idx_device = device_ndarray(out_idx) if inplace else None + out_dist_device = device_ndarray(out_dist) if inplace else None + + search_params = ivf_flat.SearchParams(n_probes=n_probes) + + ret_output = ivf_flat.search( + search_params, + index, + queries_device, + k, + neighbors=out_idx_device, + distances=out_dist_device, + ) + + if not inplace: + out_dist_device, out_idx_device = ret_output + + if not compare: + return + + out_idx = out_idx_device.copy_to_host() + out_dist = out_dist_device.copy_to_host() + + # 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_idx = nn_skl.kneighbors(queries, return_distance=False) + + recall = calc_recall(out_idx, skl_idx) + assert recall > 0.7 + + check_distances(dataset, queries, metric, out_idx, out_dist) + + +@pytest.mark.parametrize("inplace", [True, False]) +@pytest.mark.parametrize("n_rows", [10000]) +@pytest.mark.parametrize("n_cols", [10]) +@pytest.mark.parametrize("n_queries", [100]) +@pytest.mark.parametrize("n_lists", [100]) +@pytest.mark.parametrize("dtype", [np.float32, np.int8, np.uint8]) +@pytest.mark.parametrize("array_type", ["device"]) +def test_ivf_pq_dtypes( + n_rows, n_cols, n_queries, n_lists, dtype, inplace, array_type +): + # Note that inner_product tests use normalized input which we cannot + # represent in int8, therefore we test only sqeuclidean metric here. + run_ivf_flat_build_search_test( + n_rows=n_rows, + n_cols=n_cols, + n_queries=n_queries, + k=10, + n_lists=n_lists, + metric="sqeuclidean", + dtype=dtype, + inplace=inplace, + array_type=array_type, + ) + + +@pytest.mark.parametrize( + "params", + [ + pytest.param( + { + "n_rows": 0, + "n_cols": 10, + "n_queries": 10, + "k": 1, + "n_lists": 10, + }, + marks=pytest.mark.xfail(reason="empty dataset"), + ), + {"n_rows": 1, "n_cols": 10, "n_queries": 10, "k": 1, "n_lists": 1}, + {"n_rows": 10, "n_cols": 1, "n_queries": 10, "k": 10, "n_lists": 10}, + # {"n_rows": 999, "n_cols": 42, "n_queries": 453, "k": 137, + # "n_lists": 53}, + ], +) +def test_ivf_flat_n(params): + # We do not test recall, just confirm that we can handle edge cases for + # certain parameters + run_ivf_flat_build_search_test( + n_rows=params["n_rows"], + n_cols=params["n_cols"], + n_queries=params["n_queries"], + k=params["k"], + n_lists=params["n_lists"], + metric="sqeuclidean", + dtype=np.float32, + compare=False, + ) + + +@pytest.mark.parametrize( + "metric", ["sqeuclidean", "inner_product", "euclidean"] +) +@pytest.mark.parametrize("dtype", [np.float32]) +def test_ivf_flat_build_params(metric, dtype): + run_ivf_flat_build_search_test( + n_rows=10000, + n_cols=10, + n_queries=1000, + k=10, + n_lists=100, + metric=metric, + dtype=dtype, + add_data_on_build=True, + n_probes=100, + ) + + +@pytest.mark.parametrize( + "params", + [ + { + "n_lists": 100, + "trainset_fraction": 0.9, + "n_iters": 30, + }, + ], +) +def test_ivf_flat_params(params): + run_ivf_flat_build_search_test( + n_rows=10000, + n_cols=16, + n_queries=1000, + k=10, + n_lists=params["n_lists"], + metric="sqeuclidean", + dtype=np.float32, + kmeans_trainset_fraction=params.get("trainset_fraction", 1.0), + kmeans_n_iters=params.get("n_iters", 20), + ) + + +@pytest.mark.parametrize( + "params", + [ + { + "k": 10, + "n_probes": 100, + }, + { + "k": 10, + "n_probes": 99, + }, + { + "k": 10, + "n_probes": 100, + }, + { + "k": 129, + "n_probes": 100, + }, + ], +) +def test_ivf_pq_search_params(params): + run_ivf_flat_build_search_test( + n_rows=10000, + n_cols=16, + n_queries=1000, + k=params["k"], + n_lists=100, + n_probes=params["n_probes"], + metric="sqeuclidean", + dtype=np.float32, + ) + + +@pytest.mark.parametrize("dtype", [np.float32, np.int8, np.uint8]) +@pytest.mark.parametrize("array_type", ["device"]) +def test_extend(dtype, array_type): + run_ivf_flat_build_search_test( + n_rows=10000, + n_cols=10, + n_queries=100, + k=10, + n_lists=100, + metric="sqeuclidean", + dtype=dtype, + add_data_on_build=False, + array_type=array_type, + ) + + +def test_build_assertions(): + with pytest.raises(TypeError): + run_ivf_flat_build_search_test( + n_rows=1000, + n_cols=10, + n_queries=100, + k=10, + n_lists=100, + metric="sqeuclidean", + dtype=np.float64, + ) + + n_rows = 1000 + n_cols = 100 + n_queries = 212 + k = 10 + dataset = generate_data((n_rows, n_cols), np.float32) + dataset_device = device_ndarray(dataset) + + index_params = ivf_flat.IndexParams( + n_lists=50, + metric="sqeuclidean", + kmeans_n_iters=20, + kmeans_trainset_fraction=1, + add_data_on_build=False, + ) + + index = ivf_flat.Index() + + queries = generate_data((n_queries, n_cols), np.float32) + out_idx = np.zeros((n_queries, k), dtype=np.uint64) + out_dist = np.zeros((n_queries, k), dtype=np.float32) + + queries_device = device_ndarray(queries) + out_idx_device = device_ndarray(out_idx) + out_dist_device = device_ndarray(out_dist) + + search_params = ivf_flat.SearchParams(n_probes=50) + + with pytest.raises(ValueError): + # Index must be built before search + ivf_flat.search( + search_params, + index, + queries_device, + k, + out_idx_device, + out_dist_device, + ) + + index = ivf_flat.build(index_params, dataset_device) + assert index.trained + + indices = np.arange(n_rows + 1, dtype=np.uint64) + indices_device = device_ndarray(indices) + + with pytest.raises(ValueError): + # Dataset dimension mismatch + ivf_flat.extend(index, queries_device, indices_device) + + with pytest.raises(ValueError): + # indices dimension mismatch + ivf_flat.extend(index, dataset_device, indices_device) + + +@pytest.mark.parametrize( + "params", + [ + {"q_dt": np.float64}, + {"q_order": "F"}, + {"q_cols": 101}, + {"idx_dt": np.uint32}, + {"idx_order": "F"}, + {"idx_rows": 42}, + {"idx_cols": 137}, + {"dist_dt": np.float64}, + {"dist_order": "F"}, + {"dist_rows": 42}, + {"dist_cols": 137}, + ], +) +def test_search_inputs(params): + """Test with invalid input dtype, order, or dimension.""" + n_rows = 1000 + n_cols = 100 + n_queries = 256 + k = 10 + dtype = np.float32 + + q_dt = params.get("q_dt", np.float32) + q_order = params.get("q_order", "C") + queries = generate_data( + (n_queries, params.get("q_cols", n_cols)), q_dt + ).astype(q_dt, order=q_order) + queries_device = device_ndarray(queries) + + idx_dt = params.get("idx_dt", np.uint64) + idx_order = params.get("idx_order", "C") + out_idx = np.zeros( + (params.get("idx_rows", n_queries), params.get("idx_cols", k)), + dtype=idx_dt, + order=idx_order, + ) + out_idx_device = device_ndarray(out_idx) + + dist_dt = params.get("dist_dt", np.float32) + dist_order = params.get("dist_order", "C") + out_dist = np.zeros( + (params.get("dist_rows", n_queries), params.get("dist_cols", k)), + dtype=dist_dt, + order=dist_order, + ) + out_dist_device = device_ndarray(out_dist) + + index_params = ivf_flat.IndexParams( + n_lists=50, metric="sqeuclidean", add_data_on_build=True + ) + + dataset = generate_data((n_rows, n_cols), dtype) + dataset_device = device_ndarray(dataset) + index = ivf_flat.build(index_params, dataset_device) + assert index.trained + + with pytest.raises(Exception): + search_params = ivf_flat.SearchParams(n_probes=50) + ivf_flat.search( + search_params, + index, + queries_device, + k, + out_idx_device, + out_dist_device, + ) From 8faadf5316295bdadd483b0cf9da3a7e839ab78e Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 9 Mar 2023 10:03:54 -0800 Subject: [PATCH 13/20] remove references to load/save --- .../raft_runtime/neighbors/ivf_flat.hpp | 28 ----- cpp/src/distance/neighbors/ivf_flat_build.cu | 14 --- .../neighbors/ivf_flat/cpp/c_ivf_flat.pxd | 8 -- .../pylibraft/neighbors/ivf_flat/ivf_flat.pyx | 109 ------------------ 4 files changed, 159 deletions(-) diff --git a/cpp/include/raft_runtime/neighbors/ivf_flat.hpp b/cpp/include/raft_runtime/neighbors/ivf_flat.hpp index c4c751b182..a44d5e84a5 100644 --- a/cpp/include/raft_runtime/neighbors/ivf_flat.hpp +++ b/cpp/include/raft_runtime/neighbors/ivf_flat.hpp @@ -65,32 +65,4 @@ RAFT_INST_SEARCH(uint8_t, uint64_t); #undef RAFT_INST_SEARCH -// /** -// * Save the index to file. -// * -// * Experimental, both the API and the serialization format are subject to change. -// * -// * @param[in] handle the raft handle -// * @param[in] filename the filename for saving the index -// * @param[in] index IVF-PQ index -// * -// */ -// void save(raft::device_resources const& handle, -// const std::string& filename, -// const raft::neighbors::ivf_flat::index& index); - -// /** -// * Load index from file. -// * -// * Experimental, both the API and the serialization format are subject to change. -// * -// * @param[in] handle the raft handle -// * @param[in] filename the name of the file that stores the index -// * @param[in] index IVF-PQ index -// * -// */ -// void load(raft::device_resources const& handle, -// const std::string& filename, -// raft::neighbors::ivf_flat::index* index); - } // namespace raft::runtime::neighbors::ivf_flat diff --git a/cpp/src/distance/neighbors/ivf_flat_build.cu b/cpp/src/distance/neighbors/ivf_flat_build.cu index 02ab0fe653..7c8b868dad 100644 --- a/cpp/src/distance/neighbors/ivf_flat_build.cu +++ b/cpp/src/distance/neighbors/ivf_flat_build.cu @@ -59,18 +59,4 @@ RAFT_INST_BUILD_EXTEND(uint8_t, uint64_t); #undef RAFT_INST_BUILD_EXTEND -// void save(raft::device_resources const& handle, -// const std::string& filename, -// const raft::neighbors::ivf_flat::index& index) -// { -// raft::spatial::knn::ivf_flat::detail::save(handle, filename, index); -// }; - -// void load(raft::device_resources const& handle, -// const std::string& filename, -// raft::neighbors::ivf_flat::index* index) -// { -// if (!index) { RAFT_FAIL("Invalid index pointer"); } -// *index = raft::spatial::knn::ivf_flat::detail::load(handle, filename); -// }; } // namespace raft::runtime::neighbors::ivf_flat diff --git a/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/c_ivf_flat.pxd b/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/c_ivf_flat.pxd index 4e62d42cd3..982a871269 100644 --- a/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/c_ivf_flat.pxd +++ b/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/c_ivf_flat.pxd @@ -139,11 +139,3 @@ cdef extern from "raft_runtime/neighbors/ivf_flat.hpp" \ device_matrix_view[uint64_t, uint64_t, row_major] neighbors, device_matrix_view[float, uint64_t, row_major] distances, const search_params& params) except + - - # cdef void save(const device_resources& handle, - # const string& filename, - # const index[uint64_t]& index) except + - - # cdef void load(const device_resources& handle, - # const string& filename, - # index[uint64_t]* index) except + diff --git a/python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx b/python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx index ea0618fa4a..0e82070a79 100644 --- a/python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx +++ b/python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx @@ -712,112 +712,3 @@ def search(SearchParams search_params, raise ValueError("query dtype %s not supported" % queries_dt) return (distances, neighbors) - - -# @auto_sync_handle -# def save(filename, Index index, handle=None): -# """ -# Saves the index to file. - -# Saving / loading the index is experimental. The serialization format is -# subject to change. - -# Parameters -# ---------- -# filename : string -# Name of the file. -# index : Index -# Trained IVF-PQ index. -# {handle_docstring} - -# Examples -# -------- -# >>> import cupy as cp - -# >>> from pylibraft.common import DeviceResources -# >>> from pylibraft.neighbors import ivf_flat - -# >>> n_samples = 50000 -# >>> n_features = 50 -# >>> dataset = cp.random.random_sample((n_samples, n_features), -# ... dtype=cp.float32) - -# >>> # Build index -# >>> handle = DeviceResources() -# >>> index = ivf_flat.build(ivf_flat.IndexParams(), dataset, -# handle=handle) -# >>> ivf_flat.save("my_index.bin", index, handle=handle) -# """ -# if not index.trained: -# raise ValueError("Index need to be built before saving it.") - -# if handle is None: -# handle = DeviceResources() -# cdef device_resources* handle_ = \ -# handle.getHandle() - -# cdef string c_filename = filename.encode('utf-8') - -# c_ivf_flat.save(deref(handle_), c_filename, deref(index.index)) - - -# @auto_sync_handle -# def load(filename, handle=None): -# """ -# Loads index from file. - -# 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. -# {handle_docstring} - -# Returns -# ------- -# index : Index - -# Examples -# -------- -# >>> import cupy as cp - -# >>> from pylibraft.common import DeviceResources -# >>> from pylibraft.neighbors import ivf_flat - -# >>> n_samples = 50000 -# >>> n_features = 50 -# >>> dataset = cp.random.random_sample((n_samples, n_features), -# ... dtype=cp.float32) - -# >>> # Build and save index -# >>> handle = DeviceResources() -# >>> index = ivf_flat.build(ivf_flat.IndexParams(), dataset, -# handle=handle) -# >>> ivf_flat.save("my_index.bin", index, handle=handle) -# >>> del index - -# >>> n_queries = 100 -# >>> queries = cp.random.random_sample((n_queries, n_features), -# ... dtype=cp.float32) -# >>> handle = DeviceResources() -# >>> index = ivf_flat.load("my_index.bin", handle=handle) - -# >>> distances, neighbors = ivf_flat.search( -# ivf_flat.SearchParams(), index, -# ... queries, k=10, handle=handle) -# """ -# if handle is None: -# handle = DeviceResources() -# cdef device_resources* handle_ = \ -# handle.getHandle() - -# cdef string c_filename = filename.encode('utf-8') -# index = Index() - -# c_ivf_flat.load(deref(handle_), c_filename, index.index) -# index.trained = True - -# return index From 0a960b8438a2f61ec5512b65f5416858fee2fe65 Mon Sep 17 00:00:00 2001 From: divyegala Date: Fri, 10 Mar 2023 06:42:57 -0800 Subject: [PATCH 14/20] change uint64_t to int64_t --- cpp/CMakeLists.txt | 18 +++--- .../raft_runtime/neighbors/ivf_flat.hpp | 16 ++--- cpp/src/distance/neighbors/ivf_flat_build.cu | 10 +-- cpp/src/distance/neighbors/ivf_flat_search.cu | 6 +- ...64_t.cu => ivfflat_build_float_int64_t.cu} | 20 +++--- ...4_t.cu => ivfflat_build_int8_t_int64_t.cu} | 20 +++--- ..._t.cu => ivfflat_build_uint8_t_int64_t.cu} | 20 +++--- ...4_t.cu => ivfflat_extend_float_int64_t.cu} | 2 +- ..._t.cu => ivfflat_extend_int8_t_int64_t.cu} | 2 +- ...t.cu => ivfflat_extend_uint8_t_int64_t.cu} | 2 +- ...4_t.cu => ivfflat_search_float_int64_t.cu} | 2 +- ..._t.cu => ivfflat_search_int8_t_int64_t.cu} | 2 +- ...t.cu => ivfflat_search_uint8_t_int64_t.cu} | 2 +- .../neighbors/ivf_flat/cpp/c_ivf_flat.pxd | 63 +++++++++---------- .../pylibraft/neighbors/ivf_flat/ivf_flat.pyx | 59 ++++++++--------- 15 files changed, 115 insertions(+), 129 deletions(-) rename cpp/src/distance/neighbors/specializations/{ivfflat_build_float_uint64_t.cu => ivfflat_build_float_int64_t.cu} (76%) rename cpp/src/distance/neighbors/specializations/{ivfflat_build_int8_t_uint64_t.cu => ivfflat_build_int8_t_int64_t.cu} (76%) rename cpp/src/distance/neighbors/specializations/{ivfflat_build_uint8_t_uint64_t.cu => ivfflat_build_uint8_t_int64_t.cu} (76%) rename cpp/src/distance/neighbors/specializations/{ivfflat_extend_float_uint64_t.cu => ivfflat_extend_float_int64_t.cu} (98%) rename cpp/src/distance/neighbors/specializations/{ivfflat_extend_int8_t_uint64_t.cu => ivfflat_extend_int8_t_int64_t.cu} (97%) rename cpp/src/distance/neighbors/specializations/{ivfflat_extend_uint8_t_uint64_t.cu => ivfflat_extend_uint8_t_int64_t.cu} (97%) rename cpp/src/distance/neighbors/specializations/{ivfflat_search_float_uint64_t.cu => ivfflat_search_float_int64_t.cu} (97%) rename cpp/src/distance/neighbors/specializations/{ivfflat_search_int8_t_uint64_t.cu => ivfflat_search_int8_t_int64_t.cu} (97%) rename cpp/src/distance/neighbors/specializations/{ivfflat_search_uint8_t_uint64_t.cu => ivfflat_search_uint8_t_int64_t.cu} (97%) diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 5bb4f86bd2..e399704810 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -372,15 +372,15 @@ if(RAFT_COMPILE_DIST_LIBRARY) src/distance/neighbors/ivfpq_search_uint8_t_uint64_t.cu src/distance/neighbors/ivf_flat_search.cu src/distance/neighbors/ivf_flat_build.cu - src/distance/neighbors/specializations/ivfflat_build_float_uint64_t.cu - src/distance/neighbors/specializations/ivfflat_build_int8_t_uint64_t.cu - src/distance/neighbors/specializations/ivfflat_build_uint8_t_uint64_t.cu - src/distance/neighbors/specializations/ivfflat_extend_float_uint64_t.cu - src/distance/neighbors/specializations/ivfflat_extend_int8_t_uint64_t.cu - src/distance/neighbors/specializations/ivfflat_extend_uint8_t_uint64_t.cu - src/distance/neighbors/specializations/ivfflat_search_float_uint64_t.cu - src/distance/neighbors/specializations/ivfflat_search_int8_t_uint64_t.cu - src/distance/neighbors/specializations/ivfflat_search_uint8_t_uint64_t.cu + src/distance/neighbors/specializations/ivfflat_build_float_int64_t.cu + src/distance/neighbors/specializations/ivfflat_build_int8_t_int64_t.cu + src/distance/neighbors/specializations/ivfflat_build_uint8_t_int64_t.cu + src/distance/neighbors/specializations/ivfflat_extend_float_int64_t.cu + src/distance/neighbors/specializations/ivfflat_extend_int8_t_int64_t.cu + src/distance/neighbors/specializations/ivfflat_extend_uint8_t_int64_t.cu + src/distance/neighbors/specializations/ivfflat_search_float_int64_t.cu + src/distance/neighbors/specializations/ivfflat_search_int8_t_int64_t.cu + src/distance/neighbors/specializations/ivfflat_search_uint8_t_int64_t.cu src/distance/neighbors/specializations/ivfpq_build_float_uint64_t.cu src/distance/neighbors/specializations/ivfpq_build_int8_t_uint64_t.cu src/distance/neighbors/specializations/ivfpq_build_uint8_t_uint64_t.cu diff --git a/cpp/include/raft_runtime/neighbors/ivf_flat.hpp b/cpp/include/raft_runtime/neighbors/ivf_flat.hpp index a44d5e84a5..b78ab31b3d 100644 --- a/cpp/include/raft_runtime/neighbors/ivf_flat.hpp +++ b/cpp/include/raft_runtime/neighbors/ivf_flat.hpp @@ -25,7 +25,7 @@ namespace raft::runtime::neighbors::ivf_flat { // constructor. #define RAFT_INST_BUILD_EXTEND(T, IdxT) \ auto build(raft::device_resources const& handle, \ - raft::device_matrix_view dataset, \ + raft::device_matrix_view dataset, \ const raft::neighbors::ivf_flat::index_params& params) \ ->raft::neighbors::ivf_flat::index; \ \ @@ -36,7 +36,7 @@ namespace raft::runtime::neighbors::ivf_flat { ->raft::neighbors::ivf_flat::index; \ \ void build(raft::device_resources const& handle, \ - raft::device_matrix_view dataset, \ + raft::device_matrix_view dataset, \ const raft::neighbors::ivf_flat::index_params& params, \ raft::neighbors::ivf_flat::index* idx); \ \ @@ -45,9 +45,9 @@ namespace raft::runtime::neighbors::ivf_flat { raft::device_matrix_view new_vectors, \ std::optional> new_indices); -RAFT_INST_BUILD_EXTEND(float, uint64_t) -RAFT_INST_BUILD_EXTEND(int8_t, uint64_t) -RAFT_INST_BUILD_EXTEND(uint8_t, uint64_t) +RAFT_INST_BUILD_EXTEND(float, int64_t) +RAFT_INST_BUILD_EXTEND(int8_t, int64_t) +RAFT_INST_BUILD_EXTEND(uint8_t, int64_t) #undef RAFT_INST_BUILD_EXTEND @@ -59,9 +59,9 @@ RAFT_INST_BUILD_EXTEND(uint8_t, uint64_t) raft::device_matrix_view, \ raft::neighbors::ivf_flat::search_params const&); -RAFT_INST_SEARCH(float, uint64_t); -RAFT_INST_SEARCH(int8_t, uint64_t); -RAFT_INST_SEARCH(uint8_t, uint64_t); +RAFT_INST_SEARCH(float, int64_t); +RAFT_INST_SEARCH(int8_t, int64_t); +RAFT_INST_SEARCH(uint8_t, int64_t); #undef RAFT_INST_SEARCH diff --git a/cpp/src/distance/neighbors/ivf_flat_build.cu b/cpp/src/distance/neighbors/ivf_flat_build.cu index 7c8b868dad..05f4eb1324 100644 --- a/cpp/src/distance/neighbors/ivf_flat_build.cu +++ b/cpp/src/distance/neighbors/ivf_flat_build.cu @@ -21,7 +21,7 @@ namespace raft::runtime::neighbors::ivf_flat { #define RAFT_INST_BUILD_EXTEND(T, IdxT) \ auto build(raft::device_resources const& handle, \ - raft::device_matrix_view dataset, \ + raft::device_matrix_view dataset, \ const raft::neighbors::ivf_flat::index_params& params) \ ->raft::neighbors::ivf_flat::index \ { \ @@ -38,7 +38,7 @@ namespace raft::runtime::neighbors::ivf_flat { } \ \ void build(raft::device_resources const& handle, \ - raft::device_matrix_view dataset, \ + raft::device_matrix_view dataset, \ const raft::neighbors::ivf_flat::index_params& params, \ raft::neighbors::ivf_flat::index* idx) \ { \ @@ -53,9 +53,9 @@ namespace raft::runtime::neighbors::ivf_flat { raft::neighbors::ivf_flat::extend(handle, idx, new_vectors, new_indices); \ } -RAFT_INST_BUILD_EXTEND(float, uint64_t); -RAFT_INST_BUILD_EXTEND(int8_t, uint64_t); -RAFT_INST_BUILD_EXTEND(uint8_t, uint64_t); +RAFT_INST_BUILD_EXTEND(float, int64_t); +RAFT_INST_BUILD_EXTEND(int8_t, int64_t); +RAFT_INST_BUILD_EXTEND(uint8_t, int64_t); #undef RAFT_INST_BUILD_EXTEND diff --git a/cpp/src/distance/neighbors/ivf_flat_search.cu b/cpp/src/distance/neighbors/ivf_flat_search.cu index dc060460be..d046e8fbbc 100644 --- a/cpp/src/distance/neighbors/ivf_flat_search.cu +++ b/cpp/src/distance/neighbors/ivf_flat_search.cu @@ -31,9 +31,9 @@ namespace raft::runtime::neighbors::ivf_flat { handle, index, queries, neighbors, distances, params); \ } -RAFT_INST_SEARCH(float, uint64_t); -RAFT_INST_SEARCH(int8_t, uint64_t); -RAFT_INST_SEARCH(uint8_t, uint64_t); +RAFT_INST_SEARCH(float, int64_t); +RAFT_INST_SEARCH(int8_t, int64_t); +RAFT_INST_SEARCH(uint8_t, int64_t); #undef RAFT_INST_SEARCH diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_build_float_uint64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_build_float_int64_t.cu similarity index 76% rename from cpp/src/distance/neighbors/specializations/ivfflat_build_float_uint64_t.cu rename to cpp/src/distance/neighbors/specializations/ivfflat_build_float_int64_t.cu index 3cb561fdc0..0445bcfede 100644 --- a/cpp/src/distance/neighbors/specializations/ivfflat_build_float_uint64_t.cu +++ b/cpp/src/distance/neighbors/specializations/ivfflat_build_float_int64_t.cu @@ -18,18 +18,18 @@ namespace raft::neighbors::ivf_flat { -#define RAFT_MAKE_INSTANCE(T, IdxT) \ - template auto build(raft::device_resources const& handle, \ - raft::device_matrix_view dataset, \ - const index_params& params) \ - ->index; \ - \ - template void build(raft::device_resources const& handle, \ - raft::device_matrix_view dataset, \ - const index_params& params, \ +#define RAFT_MAKE_INSTANCE(T, IdxT) \ + template auto build(raft::device_resources const& handle, \ + raft::device_matrix_view dataset, \ + const index_params& params) \ + ->index; \ + \ + template void build(raft::device_resources const& handle, \ + raft::device_matrix_view dataset, \ + const index_params& params, \ index* idx); -RAFT_MAKE_INSTANCE(float, uint64_t); +RAFT_MAKE_INSTANCE(float, int64_t); #undef RAFT_MAKE_INSTANCE diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_build_int8_t_uint64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_build_int8_t_int64_t.cu similarity index 76% rename from cpp/src/distance/neighbors/specializations/ivfflat_build_int8_t_uint64_t.cu rename to cpp/src/distance/neighbors/specializations/ivfflat_build_int8_t_int64_t.cu index babca2a32c..cc58efb842 100644 --- a/cpp/src/distance/neighbors/specializations/ivfflat_build_int8_t_uint64_t.cu +++ b/cpp/src/distance/neighbors/specializations/ivfflat_build_int8_t_int64_t.cu @@ -18,18 +18,18 @@ namespace raft::neighbors::ivf_flat { -#define RAFT_MAKE_INSTANCE(T, IdxT) \ - template auto build(raft::device_resources const& handle, \ - raft::device_matrix_view dataset, \ - const index_params& params) \ - ->index; \ - \ - template void build(raft::device_resources const& handle, \ - raft::device_matrix_view dataset, \ - const index_params& params, \ +#define RAFT_MAKE_INSTANCE(T, IdxT) \ + template auto build(raft::device_resources const& handle, \ + raft::device_matrix_view dataset, \ + const index_params& params) \ + ->index; \ + \ + template void build(raft::device_resources const& handle, \ + raft::device_matrix_view dataset, \ + const index_params& params, \ index* idx); -RAFT_MAKE_INSTANCE(int8_t, uint64_t); +RAFT_MAKE_INSTANCE(int8_t, int64_t); #undef RAFT_MAKE_INSTANCE diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_build_uint8_t_uint64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_build_uint8_t_int64_t.cu similarity index 76% rename from cpp/src/distance/neighbors/specializations/ivfflat_build_uint8_t_uint64_t.cu rename to cpp/src/distance/neighbors/specializations/ivfflat_build_uint8_t_int64_t.cu index c18b3dd696..f0ac6a142a 100644 --- a/cpp/src/distance/neighbors/specializations/ivfflat_build_uint8_t_uint64_t.cu +++ b/cpp/src/distance/neighbors/specializations/ivfflat_build_uint8_t_int64_t.cu @@ -18,18 +18,18 @@ namespace raft::neighbors::ivf_flat { -#define RAFT_MAKE_INSTANCE(T, IdxT) \ - template auto build(raft::device_resources const& handle, \ - raft::device_matrix_view dataset, \ - const index_params& params) \ - ->index; \ - \ - template void build(raft::device_resources const& handle, \ - raft::device_matrix_view dataset, \ - const index_params& params, \ +#define RAFT_MAKE_INSTANCE(T, IdxT) \ + template auto build(raft::device_resources const& handle, \ + raft::device_matrix_view dataset, \ + const index_params& params) \ + ->index; \ + \ + template void build(raft::device_resources const& handle, \ + raft::device_matrix_view dataset, \ + const index_params& params, \ index* idx); -RAFT_MAKE_INSTANCE(uint8_t, uint64_t); +RAFT_MAKE_INSTANCE(uint8_t, int64_t); #undef RAFT_MAKE_INSTANCE diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_extend_float_uint64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_extend_float_int64_t.cu similarity index 98% rename from cpp/src/distance/neighbors/specializations/ivfflat_extend_float_uint64_t.cu rename to cpp/src/distance/neighbors/specializations/ivfflat_extend_float_int64_t.cu index 614a7a6a10..1a4ae888b8 100644 --- a/cpp/src/distance/neighbors/specializations/ivfflat_extend_float_uint64_t.cu +++ b/cpp/src/distance/neighbors/specializations/ivfflat_extend_float_int64_t.cu @@ -30,7 +30,7 @@ namespace raft::neighbors::ivf_flat { raft::device_matrix_view new_vectors, \ std::optional> new_indices); -RAFT_MAKE_INSTANCE(float, uint64_t); +RAFT_MAKE_INSTANCE(float, int64_t); #undef RAFT_MAKE_INSTANCE diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_extend_int8_t_uint64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_extend_int8_t_int64_t.cu similarity index 97% rename from cpp/src/distance/neighbors/specializations/ivfflat_extend_int8_t_uint64_t.cu rename to cpp/src/distance/neighbors/specializations/ivfflat_extend_int8_t_int64_t.cu index fe0f63a406..f1d5af49e5 100644 --- a/cpp/src/distance/neighbors/specializations/ivfflat_extend_int8_t_uint64_t.cu +++ b/cpp/src/distance/neighbors/specializations/ivfflat_extend_int8_t_int64_t.cu @@ -30,7 +30,7 @@ namespace raft::neighbors::ivf_flat { raft::device_matrix_view new_vectors, \ std::optional> new_indices); -RAFT_MAKE_INSTANCE(int8_t, uint64_t); +RAFT_MAKE_INSTANCE(int8_t, int64_t); #undef RAFT_MAKE_INSTANCE diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_extend_uint8_t_uint64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_extend_uint8_t_int64_t.cu similarity index 97% rename from cpp/src/distance/neighbors/specializations/ivfflat_extend_uint8_t_uint64_t.cu rename to cpp/src/distance/neighbors/specializations/ivfflat_extend_uint8_t_int64_t.cu index 187d47f842..29d2050a71 100644 --- a/cpp/src/distance/neighbors/specializations/ivfflat_extend_uint8_t_uint64_t.cu +++ b/cpp/src/distance/neighbors/specializations/ivfflat_extend_uint8_t_int64_t.cu @@ -29,7 +29,7 @@ namespace raft::neighbors::ivf_flat { raft::neighbors::ivf_flat::index* idx, \ raft::device_matrix_view new_vectors, \ std::optional> new_indices); -RAFT_MAKE_INSTANCE(uint8_t, uint64_t); +RAFT_MAKE_INSTANCE(uint8_t, int64_t); #undef RAFT_MAKE_INSTANCE diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_search_float_uint64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_search_float_int64_t.cu similarity index 97% rename from cpp/src/distance/neighbors/specializations/ivfflat_search_float_uint64_t.cu rename to cpp/src/distance/neighbors/specializations/ivfflat_search_float_int64_t.cu index fcc3b120fc..88208386fd 100644 --- a/cpp/src/distance/neighbors/specializations/ivfflat_search_float_uint64_t.cu +++ b/cpp/src/distance/neighbors/specializations/ivfflat_search_float_int64_t.cu @@ -26,7 +26,7 @@ namespace raft::neighbors::ivf_flat { raft::device_matrix_view, \ raft::neighbors::ivf_flat::search_params const&); -RAFT_MAKE_INSTANCE(float, uint64_t); +RAFT_MAKE_INSTANCE(float, int64_t); #undef RAFT_MAKE_INSTANCE diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_search_int8_t_uint64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_search_int8_t_int64_t.cu similarity index 97% rename from cpp/src/distance/neighbors/specializations/ivfflat_search_int8_t_uint64_t.cu rename to cpp/src/distance/neighbors/specializations/ivfflat_search_int8_t_int64_t.cu index 12220ebb5a..42e77fd0af 100644 --- a/cpp/src/distance/neighbors/specializations/ivfflat_search_int8_t_uint64_t.cu +++ b/cpp/src/distance/neighbors/specializations/ivfflat_search_int8_t_int64_t.cu @@ -26,7 +26,7 @@ namespace raft::neighbors::ivf_flat { raft::device_matrix_view, \ raft::neighbors::ivf_flat::search_params const&); -RAFT_MAKE_INSTANCE(int8_t, uint64_t); +RAFT_MAKE_INSTANCE(int8_t, int64_t); #undef RAFT_MAKE_INSTANCE diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_search_uint8_t_uint64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_search_uint8_t_int64_t.cu similarity index 97% rename from cpp/src/distance/neighbors/specializations/ivfflat_search_uint8_t_uint64_t.cu rename to cpp/src/distance/neighbors/specializations/ivfflat_search_uint8_t_int64_t.cu index 89859137f6..6176d59638 100644 --- a/cpp/src/distance/neighbors/specializations/ivfflat_search_uint8_t_uint64_t.cu +++ b/cpp/src/distance/neighbors/specializations/ivfflat_search_uint8_t_int64_t.cu @@ -26,7 +26,7 @@ namespace raft::neighbors::ivf_flat { raft::device_matrix_view, \ raft::neighbors::ivf_flat::search_params const&); -RAFT_MAKE_INSTANCE(uint8_t, uint64_t); +RAFT_MAKE_INSTANCE(uint8_t, int64_t); #undef RAFT_MAKE_INSTANCE diff --git a/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/c_ivf_flat.pxd b/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/c_ivf_flat.pxd index 982a871269..b539cfa7f1 100644 --- a/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/c_ivf_flat.pxd +++ b/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/c_ivf_flat.pxd @@ -23,14 +23,7 @@ import numpy as np import pylibraft.common.handle from cython.operator cimport dereference as deref -from libc.stdint cimport ( - int8_t, - int64_t, - uint8_t, - uint32_t, - uint64_t, - uintptr_t, -) +from libc.stdint cimport int8_t, int64_t, uint8_t, uint32_t, uintptr_t from libcpp cimport bool, nullptr from libcpp.string cimport string @@ -84,58 +77,58 @@ cdef extern from "raft_runtime/neighbors/ivf_flat.hpp" \ namespace "raft::runtime::neighbors::ivf_flat" nogil: cdef void build(const device_resources&, - device_matrix_view[float, uint64_t, row_major] dataset, + device_matrix_view[float, int64_t, row_major] dataset, const index_params& params, - index[float, uint64_t]* index) except + + index[float, int64_t]* index) except + cdef void build(const device_resources& handle, - device_matrix_view[int8_t, uint64_t, row_major] dataset, + device_matrix_view[int8_t, int64_t, row_major] dataset, const index_params& params, - index[int8_t, uint64_t]* index) except + + index[int8_t, int64_t]* index) except + cdef void build(const device_resources& handle, - device_matrix_view[uint8_t, uint64_t, row_major] dataset, + device_matrix_view[uint8_t, int64_t, row_major] dataset, const index_params& params, - index[uint8_t, uint64_t]* index) except + + index[uint8_t, int64_t]* index) except + cdef void extend( const device_resources& handle, - index[float, uint64_t]* index, - device_matrix_view[float, uint64_t, row_major] new_vectors, - optional[device_vector_view[uint64_t, uint64_t]] new_indices) except + + index[float, int64_t]* index, + device_matrix_view[float, int64_t, row_major] new_vectors, + optional[device_vector_view[int64_t, int64_t]] new_indices) except + cdef void extend( const device_resources& handle, - index[int8_t, uint64_t]* index, - device_matrix_view[int8_t, uint64_t, row_major] new_vectors, - optional[device_vector_view[uint64_t, uint64_t]] new_indices) except + + index[int8_t, int64_t]* index, + device_matrix_view[int8_t, int64_t, row_major] new_vectors, + optional[device_vector_view[int64_t, int64_t]] new_indices) except + cdef void extend( const device_resources& handle, - index[uint8_t, uint64_t]* index, - device_matrix_view[uint8_t, uint64_t, row_major] new_vectors, - optional[device_vector_view[uint64_t, uint64_t]] new_indices) except + + index[uint8_t, int64_t]* index, + device_matrix_view[uint8_t, int64_t, row_major] new_vectors, + optional[device_vector_view[int64_t, int64_t]] new_indices) except + cdef void search( const device_resources& handle, - const index[float, uint64_t]& index, - device_matrix_view[float, uint64_t, row_major] queries, - device_matrix_view[uint64_t, uint64_t, row_major] neighbors, - device_matrix_view[float, uint64_t, row_major] distances, + const index[float, int64_t]& index, + device_matrix_view[float, int64_t, row_major] queries, + device_matrix_view[int64_t, int64_t, row_major] neighbors, + device_matrix_view[float, int64_t, row_major] distances, const search_params& params) except + cdef void search( const device_resources& handle, - const index[int8_t, uint64_t]& index, - device_matrix_view[int8_t, uint64_t, row_major] queries, - device_matrix_view[uint64_t, uint64_t, row_major] neighbors, - device_matrix_view[float, uint64_t, row_major] distances, + const index[int8_t, int64_t]& index, + device_matrix_view[int8_t, int64_t, row_major] queries, + device_matrix_view[int64_t, int64_t, row_major] neighbors, + device_matrix_view[float, int64_t, row_major] distances, const search_params& params) except + cdef void search( const device_resources& handle, - const index[uint8_t, uint64_t]& index, - device_matrix_view[uint8_t, uint64_t, row_major] queries, - device_matrix_view[uint64_t, uint64_t, row_major] neighbors, - device_matrix_view[float, uint64_t, row_major] distances, + const index[uint8_t, int64_t]& index, + device_matrix_view[uint8_t, int64_t, row_major] queries, + device_matrix_view[int64_t, int64_t, row_major] neighbors, + device_matrix_view[float, int64_t, row_major] distances, const search_params& params) except + diff --git a/python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx b/python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx index 0e82070a79..6533896ddf 100644 --- a/python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx +++ b/python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx @@ -23,14 +23,7 @@ import warnings import numpy as np from cython.operator cimport dereference as deref -from libc.stdint cimport ( - int8_t, - int64_t, - uint8_t, - uint32_t, - uint64_t, - uintptr_t, -) +from libc.stdint cimport int8_t, int64_t, uint8_t, uint32_t, uintptr_t from libcpp cimport bool, nullptr from libcpp.string cimport string @@ -71,8 +64,8 @@ from pylibraft.neighbors.common import _check_input_array, _get_metric from pylibraft.common.mdspan cimport ( get_dmv_float, get_dmv_int8, + get_dmv_int64, get_dmv_uint8, - get_dmv_uint64, ) from pylibraft.neighbors.common cimport _get_metric_string from pylibraft.neighbors.ivf_flat.cpp.c_ivf_flat cimport ( @@ -173,7 +166,7 @@ cdef class Index: cdef class IndexFloat(Index): - cdef c_ivf_flat.index[float, uint64_t] * index + cdef c_ivf_flat.index[float, int64_t] * index def __cinit__(self, handle=None): if handle is None: @@ -184,7 +177,7 @@ cdef class IndexFloat(Index): # this is to keep track of which index type is being used # We create a placeholder object. The actual parameter values do # not matter, it will be replaced with a built index object later. - self.index = new c_ivf_flat.index[float, uint64_t]( + self.index = new c_ivf_flat.index[float, int64_t]( deref(handle_), _get_metric("sqeuclidean"), 1, False, @@ -221,7 +214,7 @@ cdef class IndexFloat(Index): cdef class IndexInt8(Index): - cdef c_ivf_flat.index[int8_t, uint64_t] * index + cdef c_ivf_flat.index[int8_t, int64_t] * index def __cinit__(self, handle=None): if handle is None: @@ -232,7 +225,7 @@ cdef class IndexInt8(Index): # this is to keep track of which index type is being used # We create a placeholder object. The actual parameter values do # not matter, it will be replaced with a built index object later. - self.index = new c_ivf_flat.index[int8_t, uint64_t]( + self.index = new c_ivf_flat.index[int8_t, int64_t]( deref(handle_), _get_metric("sqeuclidean"), 1, False, @@ -269,7 +262,7 @@ cdef class IndexInt8(Index): cdef class IndexUint8(Index): - cdef c_ivf_flat.index[uint8_t, uint64_t] * index + cdef c_ivf_flat.index[uint8_t, int64_t] * index def __cinit__(self, handle=None): if handle is None: @@ -280,7 +273,7 @@ cdef class IndexUint8(Index): # this is to keep track of which index type is being used # We create a placeholder object. The actual parameter values do # not matter, it will be replaced with a built index object later. - self.index = new c_ivf_flat.index[uint8_t, uint64_t]( + self.index = new c_ivf_flat.index[uint8_t, int64_t]( deref(handle_), _get_metric("sqeuclidean"), 1, False, @@ -373,7 +366,7 @@ def build(IndexParams index_params, dataset, handle=None): _check_input_array(dataset_cai, [np.dtype('float32'), np.dtype('byte'), np.dtype('ubyte')]) - cdef uint64_t n_rows = dataset_cai.shape[0] + cdef int64_t n_rows = dataset_cai.shape[0] cdef uint32_t dim = dataset_cai.shape[1] if handle is None: @@ -432,7 +425,7 @@ def extend(Index index, new_vectors, new_indices, handle=None): new_vectors : CUDA array interface compliant matrix shape (n_samples, dim) Supported dtype [float, int8, uint8] new_indices : CUDA array interface compliant matrix shape (n_samples, dim) - Supported dtype [uint64] + Supported dtype [int64] {handle_docstring} Returns @@ -459,7 +452,7 @@ def extend(Index index, new_vectors, new_indices, handle=None): >>> n_rows = 100 >>> more_data = cp.random.random_sample((n_rows, n_features), ... dtype=cp.float32) - >>> indices = index.size + cp.arange(n_rows, dtype=cp.uint64) + >>> indices = index.size + cp.arange(n_rows, dtype=cp.int64) >>> index = ivf_flat.extend(index, more_data, indices) >>> # Search using the built index @@ -487,18 +480,18 @@ def extend(Index index, new_vectors, new_indices, handle=None): vecs_cai = cai_wrapper(new_vectors) vecs_dt = vecs_cai.dtype - cdef uint64_t n_rows = vecs_cai.shape[0] + cdef int64_t n_rows = vecs_cai.shape[0] cdef uint32_t dim = vecs_cai.shape[1] _check_input_array(vecs_cai, [np.dtype(index.active_index_type)], exp_cols=index.dim) idx_cai = cai_wrapper(new_indices) - _check_input_array(idx_cai, [np.dtype('uint64')], exp_rows=n_rows) + _check_input_array(idx_cai, [np.dtype('int64')], exp_rows=n_rows) if len(idx_cai.shape)!=1: raise ValueError("Indices array is expected to be 1D") - cdef optional[device_vector_view[uint64_t, uint64_t]] new_indices_opt + cdef optional[device_vector_view[int64_t, int64_t]] new_indices_opt cdef IndexFloat idx_float cdef IndexInt8 idx_int8 @@ -508,8 +501,8 @@ def extend(Index index, new_vectors, new_indices, handle=None): idx_float = index if idx_float.index.size() > 0: new_indices_opt = make_device_vector_view( - idx_cai.data, - idx_cai.shape[0]) + idx_cai.data, + idx_cai.shape[0]) with cuda_interruptible(): c_ivf_flat.extend(deref(handle_), idx_float.index, @@ -519,8 +512,8 @@ def extend(Index index, new_vectors, new_indices, handle=None): idx_int8 = index if idx_int8.index[0].size() > 0: new_indices_opt = make_device_vector_view( - idx_cai.data, - idx_cai.shape[0]) + idx_cai.data, + idx_cai.shape[0]) with cuda_interruptible(): c_ivf_flat.extend(deref(handle_), idx_int8.index, @@ -530,8 +523,8 @@ def extend(Index index, new_vectors, new_indices, handle=None): idx_uint8 = index if idx_uint8.index[0].size() > 0: new_indices_opt = make_device_vector_view( - idx_cai.data, - idx_cai.shape[0]) + idx_cai.data, + idx_cai.shape[0]) with cuda_interruptible(): c_ivf_flat.extend(deref(handle_), idx_uint8.index, @@ -589,7 +582,7 @@ def search(SearchParams search_params, k : int The number of neighbors. neighbors : Optional CUDA array interface compliant matrix shape - (n_queries, k), dtype uint64_t. If supplied, neighbor + (n_queries, k), dtype int64_t. If supplied, neighbor indices will be written here in-place. (default None) distances : Optional CUDA array interface compliant matrix shape (n_queries, k) If supplied, the distances to the @@ -663,10 +656,10 @@ def search(SearchParams search_params, exp_cols=index.dim) if neighbors is None: - neighbors = device_ndarray.empty((n_queries, k), dtype='uint64') + neighbors = device_ndarray.empty((n_queries, k), dtype='int64') neighbors_cai = cai_wrapper(neighbors) - _check_input_array(neighbors_cai, [np.dtype('uint64')], + _check_input_array(neighbors_cai, [np.dtype('int64')], exp_rows=n_queries, exp_cols=k) if distances is None: @@ -687,7 +680,7 @@ def search(SearchParams search_params, c_ivf_flat.search(deref(handle_), deref(idx_float.index), get_dmv_float(queries_cai, check_shape=True), - get_dmv_uint64(neighbors_cai, check_shape=True), + get_dmv_int64(neighbors_cai, check_shape=True), get_dmv_float(distances_cai, check_shape=True), params) elif queries_dt == np.byte: @@ -696,7 +689,7 @@ def search(SearchParams search_params, c_ivf_flat.search(deref(handle_), deref(idx_int8.index), get_dmv_int8(queries_cai, check_shape=True), - get_dmv_uint64(neighbors_cai, check_shape=True), + get_dmv_int64(neighbors_cai, check_shape=True), get_dmv_float(distances_cai, check_shape=True), params) elif queries_dt == np.ubyte: @@ -705,7 +698,7 @@ def search(SearchParams search_params, c_ivf_flat.search(deref(handle_), deref(idx_uint8.index), get_dmv_uint8(queries_cai, check_shape=True), - get_dmv_uint64(neighbors_cai, check_shape=True), + get_dmv_int64(neighbors_cai, check_shape=True), get_dmv_float(distances_cai, check_shape=True), params) else: From b7128425cf60f0023129f9fef1fdc64b6b6e1634 Mon Sep 17 00:00:00 2001 From: divyegala Date: Tue, 14 Mar 2023 14:36:43 -0700 Subject: [PATCH 15/20] rearrange api signature --- cpp/include/raft/neighbors/ivf_flat.cuh | 71 ++++++++++--------- .../neighbors/specializations/ivf_flat.cuh | 23 +++--- .../raft_runtime/neighbors/ivf_flat.hpp | 22 +++--- cpp/src/distance/neighbors/ivf_flat_build.cu | 24 +++---- cpp/src/distance/neighbors/ivf_flat_search.cu | 6 +- .../ivfflat_build_float_int64_t.cu | 9 +-- .../ivfflat_build_int8_t_int64_t.cu | 9 +-- .../ivfflat_build_uint8_t_int64_t.cu | 9 +-- .../ivfflat_extend_float_int64_t.cu | 8 +-- .../ivfflat_extend_int8_t_int64_t.cu | 8 +-- .../ivfflat_extend_uint8_t_int64_t.cu | 9 +-- .../ivfflat_search_float_int64_t.cu | 4 +- .../ivfflat_search_int8_t_int64_t.cu | 4 +- .../ivfflat_search_uint8_t_int64_t.cu | 4 +- cpp/test/neighbors/ann_ivf_flat.cuh | 15 ++-- .../pylibraft/pylibraft/neighbors/common.pyx | 2 +- .../neighbors/ivf_flat/cpp/c_ivf_flat.pxd | 36 +++++----- .../pylibraft/neighbors/ivf_flat/ivf_flat.pyx | 36 +++++----- .../pylibraft/pylibraft/test/test_ivf_flat.py | 12 ++-- 19 files changed, 147 insertions(+), 164 deletions(-) diff --git a/cpp/include/raft/neighbors/ivf_flat.cuh b/cpp/include/raft/neighbors/ivf_flat.cuh index 109aaee933..7440557550 100644 --- a/cpp/include/raft/neighbors/ivf_flat.cuh +++ b/cpp/include/raft/neighbors/ivf_flat.cuh @@ -102,15 +102,16 @@ auto build(raft::device_resources const& handle, * @tparam matrix_idx_t matrix indexing type * * @param[in] handle - * @param[in] dataset a device pointer to a row-major matrix [n_rows, dim] * @param[in] params configure the index building + * @param[in] dataset a device pointer to a row-major matrix [n_rows, dim] * * @return the constructed ivf-flat index */ template auto build(raft::device_resources const& handle, - raft::device_matrix_view dataset, - const index_params& params) -> index + const index_params& params, + raft::device_matrix_view dataset) + -> index { return raft::spatial::knn::ivf_flat::detail::build(handle, params, @@ -133,8 +134,8 @@ auto build(raft::device_resources const& handle, * // use default index parameters * ivf_flat::index_params index_params; * // create and fill the index from a [N, D] dataset - * ivf_flat::index *index_ptr; - * ivf_flat::build(handle, dataset, index_params, index_ptr); + * ivf_flat::index index; + * ivf_flat::build(handle, dataset, index_params, index); * // use default search parameters * ivf_flat::search_params search_params; * // search K nearest neighbours for each of the N queries @@ -147,22 +148,22 @@ auto build(raft::device_resources const& handle, * @tparam matrix_idx_t matrix indexing type * * @param[in] handle - * @param[in] dataset a device pointer to a row-major matrix [n_rows, dim] * @param[in] params configure the index building - * @param[out] idx pointer to ivf_flat::index + * @param[in] dataset raft::device_matrix_view to a row-major matrix [n_rows, dim] + * @param[out] idx reference to ivf_flat::index * */ template void build(raft::device_resources const& handle, - raft::device_matrix_view dataset, const index_params& params, - raft::neighbors::ivf_flat::index* idx) + raft::device_matrix_view dataset, + raft::neighbors::ivf_flat::index& idx) { - *idx = raft::spatial::knn::ivf_flat::detail::build(handle, - params, - dataset.data_handle(), - static_cast(dataset.extent(0)), - static_cast(dataset.extent(1))); + idx = raft::spatial::knn::ivf_flat::detail::build(handle, + params, + dataset.data_handle(), + static_cast(dataset.extent(0)), + static_cast(dataset.extent(1))); } /** @} */ @@ -240,20 +241,19 @@ auto extend(raft::device_resources const& handle, * @tparam matrix_idx_t matrix indexing type * * @param[in] handle - * @param[in] orig_index original index - * @param[in] new_vectors a device pointer to a row-major matrix [n_rows, index.dim()] - * @param[in] new_indices a device pointer to a vector of indices [n_rows]. - * If the original index is empty (`orig_index.size() == 0`), you can pass `nullptr` + * @param[in] new_vectors raft::device_matrix_view to a row-major matrix [n_rows, index.dim()] + * @param[in] new_indices optional raft::device_matrix_view to a vector of indices [n_rows]. + * If the original index is empty (`orig_index.size() == 0`), you can pass `std::nullopt` * here to imply a continuous range `[0...n_rows)`. + * @param[in] orig_index original index * * @return the constructed extended ivf-flat index */ template auto extend(raft::device_resources const& handle, - const index& orig_index, raft::device_matrix_view new_vectors, - std::optional> new_indices = std::nullopt) - -> index + std::optional> new_indices, + const index& orig_index) -> index { return extend( handle, @@ -318,7 +318,8 @@ void extend(raft::device_resources const& handle, * // train the index from a [N, D] dataset * auto index_empty = ivf_flat::build(handle, dataset, index_params, dataset); * // fill the index with the data - * ivf_flat::extend(handle, index_empty, dataset); + * std::optional> no_op = std::nullopt; + * ivf_flat::extend(handle, dataset, no_opt, &index_empty); * @endcode * * @tparam value_t data element type @@ -327,23 +328,23 @@ void extend(raft::device_resources const& handle, * @tparam matrix_idx_t matrix indexing type * * @param[in] handle - * @param[inout] index - * @param[in] new_vectors a device pointer to a row-major matrix [n_rows, index.dim()] - * @param[in] new_indices a device pointer to a vector of indices [n_rows]. + * @param[in] new_vectors raft::device_matrix_view to a row-major matrix [n_rows, index.dim()] + * @param[in] new_indices optional raft::device_matrix_view to a vector of indices [n_rows]. * If the original index is empty (`orig_index.size() == 0`), you can pass `std::nullopt` * here to imply a continuous range `[0...n_rows)`. + * @param[inout] index pointer to index, to be overwritten in-place */ template void extend(raft::device_resources const& handle, - index* index, raft::device_matrix_view new_vectors, - std::optional> new_indices = std::nullopt) + std::optional> new_indices, + index* index) { - *index = extend(handle, - *index, - new_vectors.data_handle(), - new_indices.has_value() ? new_indices.value().data_handle() : nullptr, - static_cast(new_vectors.extent(0))); + extend(handle, + index, + new_vectors.data_handle(), + new_indices.has_value() ? new_indices.value().data_handle() : nullptr, + static_cast(new_vectors.extent(0))); } /** @} */ @@ -437,20 +438,20 @@ void search(raft::device_resources const& handle, * @tparam matrix_idx_t matrix indexing type * * @param[in] handle + * @param[in] params configure the search * @param[in] index ivf-flat constructed index * @param[in] queries a device pointer to a row-major matrix [n_queries, index->dim()] * @param[out] neighbors a device pointer to the indices of the neighbors in the source dataset * [n_queries, k] * @param[out] distances a device pointer to the distances to the selected neighbors [n_queries, k] - * @param[in] params configure the search */ template void search(raft::device_resources const& handle, + const search_params& params, const index& index, raft::device_matrix_view queries, raft::device_matrix_view neighbors, - raft::device_matrix_view distances, - const search_params& params) + raft::device_matrix_view distances) { RAFT_EXPECTS( queries.extent(0) == neighbors.extent(0) && queries.extent(0) == distances.extent(0), diff --git a/cpp/include/raft/neighbors/specializations/ivf_flat.cuh b/cpp/include/raft/neighbors/specializations/ivf_flat.cuh index ed4480c7fa..02e1cbebb0 100644 --- a/cpp/include/raft/neighbors/specializations/ivf_flat.cuh +++ b/cpp/include/raft/neighbors/specializations/ivf_flat.cuh @@ -22,34 +22,29 @@ namespace raft::neighbors::ivf_flat { #define RAFT_INST(T, IdxT) \ extern template auto build(raft::device_resources const& handle, \ - raft::device_matrix_view dataset, \ - const index_params& params) \ + const index_params& params, \ + raft::device_matrix_view dataset) \ ->index; \ \ extern template auto extend( \ raft::device_resources const& handle, \ - const index& orig_index, \ raft::device_matrix_view new_vectors, \ - std::optional> new_indices) \ + std::optional> new_indices, \ + const index& orig_index) \ ->index; \ \ - extern template void build(raft::device_resources const& handle, \ - raft::device_matrix_view dataset, \ - const index_params& params, \ - index* idx); \ - \ extern template void extend( \ raft::device_resources const& handle, \ - index* idx, \ raft::device_matrix_view new_vectors, \ - std::optional> new_indices); \ + std::optional> new_indices, \ + raft::neighbors::ivf_flat::index* idx); \ \ extern template void search(raft::device_resources const&, \ - const index&, \ + raft::neighbors::ivf_flat::search_params const&, \ + const raft::neighbors::ivf_flat::index&, \ raft::device_matrix_view, \ raft::device_matrix_view, \ - raft::device_matrix_view, \ - search_params const&); + raft::device_matrix_view); RAFT_INST(float, uint64_t); RAFT_INST(int8_t, uint64_t); diff --git a/cpp/include/raft_runtime/neighbors/ivf_flat.hpp b/cpp/include/raft_runtime/neighbors/ivf_flat.hpp index b78ab31b3d..18ea064015 100644 --- a/cpp/include/raft_runtime/neighbors/ivf_flat.hpp +++ b/cpp/include/raft_runtime/neighbors/ivf_flat.hpp @@ -25,25 +25,25 @@ namespace raft::runtime::neighbors::ivf_flat { // constructor. #define RAFT_INST_BUILD_EXTEND(T, IdxT) \ auto build(raft::device_resources const& handle, \ - raft::device_matrix_view dataset, \ - const raft::neighbors::ivf_flat::index_params& params) \ + const raft::neighbors::ivf_flat::index_params& params, \ + raft::device_matrix_view dataset) \ ->raft::neighbors::ivf_flat::index; \ \ auto extend(raft::device_resources const& handle, \ - const raft::neighbors::ivf_flat::index& orig_index, \ raft::device_matrix_view new_vectors, \ - std::optional> new_indices) \ + std::optional> new_indices, \ + const raft::neighbors::ivf_flat::index& orig_index) \ ->raft::neighbors::ivf_flat::index; \ \ void build(raft::device_resources const& handle, \ - raft::device_matrix_view dataset, \ const raft::neighbors::ivf_flat::index_params& params, \ - raft::neighbors::ivf_flat::index* idx); \ + raft::device_matrix_view dataset, \ + raft::neighbors::ivf_flat::index& idx); \ \ void extend(raft::device_resources const& handle, \ - raft::neighbors::ivf_flat::index* idx, \ raft::device_matrix_view new_vectors, \ - std::optional> new_indices); + std::optional> new_indices, \ + raft::neighbors::ivf_flat::index* idx); RAFT_INST_BUILD_EXTEND(float, int64_t) RAFT_INST_BUILD_EXTEND(int8_t, int64_t) @@ -53,11 +53,11 @@ RAFT_INST_BUILD_EXTEND(uint8_t, int64_t) #define RAFT_INST_SEARCH(T, IdxT) \ void search(raft::device_resources const&, \ - const raft::neighbors::ivf_flat::index&, \ + raft::neighbors::ivf_flat::search_params const&, \ + raft::neighbors::ivf_flat::index const&, \ raft::device_matrix_view, \ raft::device_matrix_view, \ - raft::device_matrix_view, \ - raft::neighbors::ivf_flat::search_params const&); + raft::device_matrix_view); RAFT_INST_SEARCH(float, int64_t); RAFT_INST_SEARCH(int8_t, int64_t); diff --git a/cpp/src/distance/neighbors/ivf_flat_build.cu b/cpp/src/distance/neighbors/ivf_flat_build.cu index 05f4eb1324..9ceb455a9a 100644 --- a/cpp/src/distance/neighbors/ivf_flat_build.cu +++ b/cpp/src/distance/neighbors/ivf_flat_build.cu @@ -21,36 +21,36 @@ namespace raft::runtime::neighbors::ivf_flat { #define RAFT_INST_BUILD_EXTEND(T, IdxT) \ auto build(raft::device_resources const& handle, \ - raft::device_matrix_view dataset, \ - const raft::neighbors::ivf_flat::index_params& params) \ + const raft::neighbors::ivf_flat::index_params& params, \ + raft::device_matrix_view dataset) \ ->raft::neighbors::ivf_flat::index \ { \ - return raft::neighbors::ivf_flat::build(handle, dataset, params); \ + return raft::neighbors::ivf_flat::build(handle, params, dataset); \ } \ auto extend(raft::device_resources const& handle, \ - const raft::neighbors::ivf_flat::index& orig_index, \ raft::device_matrix_view new_vectors, \ - std::optional> new_indices) \ + std::optional> new_indices, \ + const raft::neighbors::ivf_flat::index& orig_index) \ ->raft::neighbors::ivf_flat::index \ { \ return raft::neighbors::ivf_flat::extend( \ - handle, orig_index, new_vectors, new_indices); \ + handle, new_vectors, new_indices, orig_index); \ } \ \ void build(raft::device_resources const& handle, \ - raft::device_matrix_view dataset, \ const raft::neighbors::ivf_flat::index_params& params, \ - raft::neighbors::ivf_flat::index* idx) \ + raft::device_matrix_view dataset, \ + raft::neighbors::ivf_flat::index& idx) \ { \ - *idx = raft::neighbors::ivf_flat::build(handle, dataset, params); \ + idx = build(handle, params, dataset); \ } \ \ void extend(raft::device_resources const& handle, \ - raft::neighbors::ivf_flat::index* idx, \ raft::device_matrix_view new_vectors, \ - std::optional> new_indices) \ + std::optional> new_indices, \ + raft::neighbors::ivf_flat::index* idx) \ { \ - raft::neighbors::ivf_flat::extend(handle, idx, new_vectors, new_indices); \ + raft::neighbors::ivf_flat::extend(handle, new_vectors, new_indices, idx); \ } RAFT_INST_BUILD_EXTEND(float, int64_t); diff --git a/cpp/src/distance/neighbors/ivf_flat_search.cu b/cpp/src/distance/neighbors/ivf_flat_search.cu index d046e8fbbc..f6d33df5be 100644 --- a/cpp/src/distance/neighbors/ivf_flat_search.cu +++ b/cpp/src/distance/neighbors/ivf_flat_search.cu @@ -21,14 +21,14 @@ namespace raft::runtime::neighbors::ivf_flat { #define RAFT_INST_SEARCH(T, IdxT) \ void search(raft::device_resources const& handle, \ + raft::neighbors::ivf_flat::search_params const& params, \ const raft::neighbors::ivf_flat::index& index, \ raft::device_matrix_view queries, \ raft::device_matrix_view neighbors, \ - raft::device_matrix_view distances, \ - raft::neighbors::ivf_flat::search_params const& params) \ + raft::device_matrix_view distances) \ { \ raft::neighbors::ivf_flat::search( \ - handle, index, queries, neighbors, distances, params); \ + handle, params, index, queries, neighbors, distances); \ } RAFT_INST_SEARCH(float, int64_t); diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_build_float_int64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_build_float_int64_t.cu index 0445bcfede..a6c0723172 100644 --- a/cpp/src/distance/neighbors/specializations/ivfflat_build_float_int64_t.cu +++ b/cpp/src/distance/neighbors/specializations/ivfflat_build_float_int64_t.cu @@ -20,14 +20,9 @@ namespace raft::neighbors::ivf_flat { #define RAFT_MAKE_INSTANCE(T, IdxT) \ template auto build(raft::device_resources const& handle, \ - raft::device_matrix_view dataset, \ - const index_params& params) \ - ->index; \ - \ - template void build(raft::device_resources const& handle, \ - raft::device_matrix_view dataset, \ const index_params& params, \ - index* idx); + raft::device_matrix_view dataset) \ + ->index; RAFT_MAKE_INSTANCE(float, int64_t); diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_build_int8_t_int64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_build_int8_t_int64_t.cu index cc58efb842..50b7d069c9 100644 --- a/cpp/src/distance/neighbors/specializations/ivfflat_build_int8_t_int64_t.cu +++ b/cpp/src/distance/neighbors/specializations/ivfflat_build_int8_t_int64_t.cu @@ -20,14 +20,9 @@ namespace raft::neighbors::ivf_flat { #define RAFT_MAKE_INSTANCE(T, IdxT) \ template auto build(raft::device_resources const& handle, \ - raft::device_matrix_view dataset, \ - const index_params& params) \ - ->index; \ - \ - template void build(raft::device_resources const& handle, \ - raft::device_matrix_view dataset, \ const index_params& params, \ - index* idx); + raft::device_matrix_view dataset) \ + ->index; RAFT_MAKE_INSTANCE(int8_t, int64_t); diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_build_uint8_t_int64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_build_uint8_t_int64_t.cu index f0ac6a142a..d2ac83e770 100644 --- a/cpp/src/distance/neighbors/specializations/ivfflat_build_uint8_t_int64_t.cu +++ b/cpp/src/distance/neighbors/specializations/ivfflat_build_uint8_t_int64_t.cu @@ -20,14 +20,9 @@ namespace raft::neighbors::ivf_flat { #define RAFT_MAKE_INSTANCE(T, IdxT) \ template auto build(raft::device_resources const& handle, \ - raft::device_matrix_view dataset, \ - const index_params& params) \ - ->index; \ - \ - template void build(raft::device_resources const& handle, \ - raft::device_matrix_view dataset, \ const index_params& params, \ - index* idx); + raft::device_matrix_view dataset) \ + ->index; RAFT_MAKE_INSTANCE(uint8_t, int64_t); diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_extend_float_int64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_extend_float_int64_t.cu index 1a4ae888b8..5af3e37acd 100644 --- a/cpp/src/distance/neighbors/specializations/ivfflat_extend_float_int64_t.cu +++ b/cpp/src/distance/neighbors/specializations/ivfflat_extend_float_int64_t.cu @@ -20,15 +20,15 @@ namespace raft::neighbors::ivf_flat { #define RAFT_MAKE_INSTANCE(T, IdxT) \ template auto extend(raft::device_resources const& handle, \ - const raft::neighbors::ivf_flat::index& orig_index, \ raft::device_matrix_view new_vectors, \ - std::optional> new_indices) \ + std::optional> new_indices, \ + const raft::neighbors::ivf_flat::index& orig_index) \ ->raft::neighbors::ivf_flat::index; \ \ template void extend(raft::device_resources const& handle, \ - raft::neighbors::ivf_flat::index* idx, \ raft::device_matrix_view new_vectors, \ - std::optional> new_indices); + std::optional> new_indices, \ + raft::neighbors::ivf_flat::index* idx); RAFT_MAKE_INSTANCE(float, int64_t); diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_extend_int8_t_int64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_extend_int8_t_int64_t.cu index f1d5af49e5..f89d523426 100644 --- a/cpp/src/distance/neighbors/specializations/ivfflat_extend_int8_t_int64_t.cu +++ b/cpp/src/distance/neighbors/specializations/ivfflat_extend_int8_t_int64_t.cu @@ -20,15 +20,15 @@ namespace raft::neighbors::ivf_flat { #define RAFT_MAKE_INSTANCE(T, IdxT) \ template auto extend(raft::device_resources const& handle, \ - const raft::neighbors::ivf_flat::index& orig_index, \ raft::device_matrix_view new_vectors, \ - std::optional> new_indices) \ + std::optional> new_indices, \ + const raft::neighbors::ivf_flat::index& orig_index) \ ->raft::neighbors::ivf_flat::index; \ \ template void extend(raft::device_resources const& handle, \ - raft::neighbors::ivf_flat::index* idx, \ raft::device_matrix_view new_vectors, \ - std::optional> new_indices); + std::optional> new_indices, \ + raft::neighbors::ivf_flat::index* idx); RAFT_MAKE_INSTANCE(int8_t, int64_t); diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_extend_uint8_t_int64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_extend_uint8_t_int64_t.cu index 29d2050a71..1fc8bca9ea 100644 --- a/cpp/src/distance/neighbors/specializations/ivfflat_extend_uint8_t_int64_t.cu +++ b/cpp/src/distance/neighbors/specializations/ivfflat_extend_uint8_t_int64_t.cu @@ -20,15 +20,16 @@ namespace raft::neighbors::ivf_flat { #define RAFT_MAKE_INSTANCE(T, IdxT) \ template auto extend(raft::device_resources const& handle, \ - const raft::neighbors::ivf_flat::index& orig_index, \ raft::device_matrix_view new_vectors, \ - std::optional> new_indices) \ + std::optional> new_indices, \ + const raft::neighbors::ivf_flat::index& orig_index) \ ->raft::neighbors::ivf_flat::index; \ \ template void extend(raft::device_resources const& handle, \ - raft::neighbors::ivf_flat::index* idx, \ raft::device_matrix_view new_vectors, \ - std::optional> new_indices); + std::optional> new_indices, \ + raft::neighbors::ivf_flat::index* idx); + RAFT_MAKE_INSTANCE(uint8_t, int64_t); #undef RAFT_MAKE_INSTANCE diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_search_float_int64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_search_float_int64_t.cu index 88208386fd..84d909d4bf 100644 --- a/cpp/src/distance/neighbors/specializations/ivfflat_search_float_int64_t.cu +++ b/cpp/src/distance/neighbors/specializations/ivfflat_search_float_int64_t.cu @@ -20,11 +20,11 @@ namespace raft::neighbors::ivf_flat { #define RAFT_MAKE_INSTANCE(T, IdxT) \ template void search(raft::device_resources const&, \ + raft::neighbors::ivf_flat::search_params const&, \ const raft::neighbors::ivf_flat::index&, \ raft::device_matrix_view, \ raft::device_matrix_view, \ - raft::device_matrix_view, \ - raft::neighbors::ivf_flat::search_params const&); + raft::device_matrix_view); RAFT_MAKE_INSTANCE(float, int64_t); diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_search_int8_t_int64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_search_int8_t_int64_t.cu index 42e77fd0af..26cc91e224 100644 --- a/cpp/src/distance/neighbors/specializations/ivfflat_search_int8_t_int64_t.cu +++ b/cpp/src/distance/neighbors/specializations/ivfflat_search_int8_t_int64_t.cu @@ -20,11 +20,11 @@ namespace raft::neighbors::ivf_flat { #define RAFT_MAKE_INSTANCE(T, IdxT) \ template void search(raft::device_resources const&, \ + raft::neighbors::ivf_flat::search_params const&, \ const raft::neighbors::ivf_flat::index&, \ raft::device_matrix_view, \ raft::device_matrix_view, \ - raft::device_matrix_view, \ - raft::neighbors::ivf_flat::search_params const&); + raft::device_matrix_view); RAFT_MAKE_INSTANCE(int8_t, int64_t); diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_search_uint8_t_int64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_search_uint8_t_int64_t.cu index 6176d59638..87de7183a0 100644 --- a/cpp/src/distance/neighbors/specializations/ivfflat_search_uint8_t_int64_t.cu +++ b/cpp/src/distance/neighbors/specializations/ivfflat_search_uint8_t_int64_t.cu @@ -20,11 +20,11 @@ namespace raft::neighbors::ivf_flat { #define RAFT_MAKE_INSTANCE(T, IdxT) \ template void search(raft::device_resources const&, \ + raft::neighbors::ivf_flat::search_params const&, \ const raft::neighbors::ivf_flat::index&, \ raft::device_matrix_view, \ raft::device_matrix_view, \ - raft::device_matrix_view, \ - raft::neighbors::ivf_flat::search_params const&); + raft::device_matrix_view); RAFT_MAKE_INSTANCE(uint8_t, int64_t); diff --git a/cpp/test/neighbors/ann_ivf_flat.cuh b/cpp/test/neighbors/ann_ivf_flat.cuh index 8d6b184c19..8495c292a3 100644 --- a/cpp/test/neighbors/ann_ivf_flat.cuh +++ b/cpp/test/neighbors/ann_ivf_flat.cuh @@ -159,7 +159,7 @@ class AnnIVFFlatTest : public ::testing::TestWithParam> { auto database_view = raft::make_device_matrix_view( (const DataT*)database.data(), ps.num_db_vecs, ps.dim); - auto index = ivf_flat::build(handle_, database_view, index_params); + auto idx = ivf_flat::build(handle_, index_params, database_view); rmm::device_uvector vector_indices(ps.num_db_vecs, stream_); thrust::sequence(handle_.get_thrust_policy(), @@ -172,7 +172,8 @@ class AnnIVFFlatTest : public ::testing::TestWithParam> { auto half_of_data_view = raft::make_device_matrix_view( (const DataT*)database.data(), half_of_data, ps.dim); - auto index_2 = ivf_flat::extend(handle_, index, half_of_data_view); + const std::optional> no_opt = std::nullopt; + index index_2 = ivf_flat::extend(handle_, half_of_data_view, no_opt, idx); auto new_half_of_data_view = raft::make_device_matrix_view( database.data() + half_of_data * ps.dim, IdxT(ps.num_db_vecs) - half_of_data, ps.dim); @@ -181,10 +182,10 @@ class AnnIVFFlatTest : public ::testing::TestWithParam> { vector_indices.data() + half_of_data, IdxT(ps.num_db_vecs) - half_of_data); ivf_flat::extend(handle_, - &index_2, new_half_of_data_view, std::make_optional>( - new_half_of_data_indices_view)); + new_half_of_data_indices_view), + index_2); auto search_queries_view = raft::make_device_matrix_view( search_queries.data(), ps.num_queries, ps.dim); @@ -198,11 +199,11 @@ class AnnIVFFlatTest : public ::testing::TestWithParam> { raft::spatial::knn::ivf_flat::detail::deserialize(handle_, "ivf_flat_index"); ivf_flat::search(handle_, + search_params, index_loaded, search_queries_view, indices_out_view, - dists_out_view, - search_params); + dists_out_view); update_host(distances_ivfflat.data(), distances_ivfflat_dev.data(), queries_size, stream_); update_host(indices_ivfflat.data(), indices_ivfflat_dev.data(), queries_size, stream_); @@ -241,7 +242,7 @@ class AnnIVFFlatTest : public ::testing::TestWithParam> { } else { // The centers must be immutable ASSERT_TRUE(raft::devArrMatch(index_2.centers().data_handle(), - index.centers().data_handle(), + idx.centers().data_handle(), index_2.centers().size(), raft::Compare(), stream_)); diff --git a/python/pylibraft/pylibraft/neighbors/common.pyx b/python/pylibraft/pylibraft/neighbors/common.pyx index 1de782c185..a8380b589b 100644 --- a/python/pylibraft/pylibraft/neighbors/common.pyx +++ b/python/pylibraft/pylibraft/neighbors/common.pyx @@ -47,7 +47,7 @@ cdef _get_metric_string(DistanceType metric): def _check_input_array(cai, exp_dt, exp_rows=None, exp_cols=None): if cai.dtype not in exp_dt: - raise TypeError("dtype %s not supported" % cai["typestr"]) + raise TypeError("dtype %s not supported" % cai.dtype) if not cai.c_contiguous: raise ValueError("Row major input is expected") diff --git a/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/c_ivf_flat.pxd b/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/c_ivf_flat.pxd index b539cfa7f1..94f7cfceb6 100644 --- a/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/c_ivf_flat.pxd +++ b/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/c_ivf_flat.pxd @@ -77,58 +77,58 @@ cdef extern from "raft_runtime/neighbors/ivf_flat.hpp" \ namespace "raft::runtime::neighbors::ivf_flat" nogil: cdef void build(const device_resources&, - device_matrix_view[float, int64_t, row_major] dataset, const index_params& params, - index[float, int64_t]* index) except + + device_matrix_view[float, int64_t, row_major] dataset, + index[float, int64_t]& index) except + cdef void build(const device_resources& handle, - device_matrix_view[int8_t, int64_t, row_major] dataset, const index_params& params, - index[int8_t, int64_t]* index) except + + device_matrix_view[int8_t, int64_t, row_major] dataset, + index[int8_t, int64_t]& index) except + cdef void build(const device_resources& handle, - device_matrix_view[uint8_t, int64_t, row_major] dataset, const index_params& params, - index[uint8_t, int64_t]* index) except + + device_matrix_view[uint8_t, int64_t, row_major] dataset, + index[uint8_t, int64_t]& index) except + cdef void extend( const device_resources& handle, - index[float, int64_t]* index, device_matrix_view[float, int64_t, row_major] new_vectors, - optional[device_vector_view[int64_t, int64_t]] new_indices) except + + optional[device_vector_view[int64_t, int64_t]] new_indices, + index[float, int64_t]* index) except + cdef void extend( const device_resources& handle, - index[int8_t, int64_t]* index, device_matrix_view[int8_t, int64_t, row_major] new_vectors, - optional[device_vector_view[int64_t, int64_t]] new_indices) except + + optional[device_vector_view[int64_t, int64_t]] new_indices, + index[int8_t, int64_t]* index) except + cdef void extend( const device_resources& handle, - index[uint8_t, int64_t]* index, device_matrix_view[uint8_t, int64_t, row_major] new_vectors, - optional[device_vector_view[int64_t, int64_t]] new_indices) except + + optional[device_vector_view[int64_t, int64_t]] new_indices, + index[uint8_t, int64_t]* index) except + cdef void search( const device_resources& handle, + const search_params& params, const index[float, int64_t]& index, device_matrix_view[float, int64_t, row_major] queries, device_matrix_view[int64_t, int64_t, row_major] neighbors, - device_matrix_view[float, int64_t, row_major] distances, - const search_params& params) except + + device_matrix_view[float, int64_t, row_major] distances) except + cdef void search( const device_resources& handle, + const search_params& params, const index[int8_t, int64_t]& index, device_matrix_view[int8_t, int64_t, row_major] queries, device_matrix_view[int64_t, int64_t, row_major] neighbors, - device_matrix_view[float, int64_t, row_major] distances, - const search_params& params) except + + device_matrix_view[float, int64_t, row_major] distances) except + cdef void search( const device_resources& handle, + const search_params& params, const index[uint8_t, int64_t]& index, device_matrix_view[uint8_t, int64_t, row_major] queries, device_matrix_view[int64_t, int64_t, row_major] neighbors, - device_matrix_view[float, int64_t, row_major] distances, - const search_params& params) except + + device_matrix_view[float, int64_t, row_major] distances) except + diff --git a/python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx b/python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx index 6533896ddf..b89ac3ada5 100644 --- a/python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx +++ b/python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx @@ -383,9 +383,9 @@ def build(IndexParams index_params, dataset, handle=None): idx_float.active_index_type = "float32" with cuda_interruptible(): c_ivf_flat.build(deref(handle_), - get_dmv_float(dataset_cai, check_shape=True), index_params.params, - idx_float.index) + get_dmv_float(dataset_cai, check_shape=True), + deref(idx_float.index)) idx_float.trained = True return idx_float elif dataset_dt == np.byte: @@ -393,9 +393,9 @@ def build(IndexParams index_params, dataset, handle=None): idx_int8.active_index_type = "byte" with cuda_interruptible(): c_ivf_flat.build(deref(handle_), - get_dmv_int8(dataset_cai, check_shape=True), index_params.params, - idx_int8.index) + get_dmv_int8(dataset_cai, check_shape=True), + deref(idx_int8.index)) idx_int8.trained = True return idx_int8 elif dataset_dt == np.ubyte: @@ -403,9 +403,9 @@ def build(IndexParams index_params, dataset, handle=None): idx_uint8.active_index_type = "ubyte" with cuda_interruptible(): c_ivf_flat.build(deref(handle_), - get_dmv_uint8(dataset_cai, check_shape=True), index_params.params, - idx_uint8.index) + get_dmv_uint8(dataset_cai, check_shape=True), + deref(idx_uint8.index)) idx_uint8.trained = True return idx_uint8 else: @@ -505,9 +505,9 @@ def extend(Index index, new_vectors, new_indices, handle=None): idx_cai.shape[0]) with cuda_interruptible(): c_ivf_flat.extend(deref(handle_), - idx_float.index, get_dmv_float(vecs_cai, check_shape=True), - new_indices_opt) + new_indices_opt, + idx_float.index) elif vecs_dt == np.int8: idx_int8 = index if idx_int8.index[0].size() > 0: @@ -516,9 +516,9 @@ def extend(Index index, new_vectors, new_indices, handle=None): idx_cai.shape[0]) with cuda_interruptible(): c_ivf_flat.extend(deref(handle_), - idx_int8.index, get_dmv_int8(vecs_cai, check_shape=True), - new_indices_opt) + new_indices_opt, + idx_int8.index) elif vecs_dt == np.uint8: idx_uint8 = index if idx_uint8.index[0].size() > 0: @@ -527,9 +527,9 @@ def extend(Index index, new_vectors, new_indices, handle=None): idx_cai.shape[0]) with cuda_interruptible(): c_ivf_flat.extend(deref(handle_), - idx_uint8.index, get_dmv_uint8(vecs_cai, check_shape=True), - new_indices_opt) + new_indices_opt, + idx_uint8.index) else: raise TypeError("query dtype %s not supported" % vecs_dt) @@ -678,29 +678,29 @@ def search(SearchParams search_params, idx_float = index with cuda_interruptible(): c_ivf_flat.search(deref(handle_), + params, deref(idx_float.index), get_dmv_float(queries_cai, check_shape=True), get_dmv_int64(neighbors_cai, check_shape=True), - get_dmv_float(distances_cai, check_shape=True), - params) + get_dmv_float(distances_cai, check_shape=True)) elif queries_dt == np.byte: idx_int8 = index with cuda_interruptible(): c_ivf_flat.search(deref(handle_), + params, deref(idx_int8.index), get_dmv_int8(queries_cai, check_shape=True), get_dmv_int64(neighbors_cai, check_shape=True), - get_dmv_float(distances_cai, check_shape=True), - params) + get_dmv_float(distances_cai, check_shape=True)) elif queries_dt == np.ubyte: idx_uint8 = index with cuda_interruptible(): c_ivf_flat.search(deref(handle_), + params, deref(idx_uint8.index), get_dmv_uint8(queries_cai, check_shape=True), get_dmv_int64(neighbors_cai, check_shape=True), - get_dmv_float(distances_cai, check_shape=True), - params) + get_dmv_float(distances_cai, check_shape=True)) else: raise ValueError("query dtype %s not supported" % queries_dt) diff --git a/python/pylibraft/pylibraft/test/test_ivf_flat.py b/python/pylibraft/pylibraft/test/test_ivf_flat.py index ce6d3c6a4b..593980f7c8 100644 --- a/python/pylibraft/pylibraft/test/test_ivf_flat.py +++ b/python/pylibraft/pylibraft/test/test_ivf_flat.py @@ -116,8 +116,8 @@ def run_ivf_flat_build_search_test( if not add_data_on_build: dataset_1 = dataset[: n_rows // 2, :] dataset_2 = dataset[n_rows // 2 :, :] - indices_1 = np.arange(n_rows // 2, dtype=np.uint64) - indices_2 = np.arange(n_rows // 2, n_rows, dtype=np.uint64) + indices_1 = np.arange(n_rows // 2, dtype=np.int64) + indices_2 = np.arange(n_rows // 2, n_rows, dtype=np.int64) if array_type == "device": dataset_1_device = device_ndarray(dataset_1) dataset_2_device = device_ndarray(dataset_2) @@ -132,7 +132,7 @@ def run_ivf_flat_build_search_test( assert index.size >= n_rows queries = generate_data((n_queries, n_cols), dtype) - out_idx = np.zeros((n_queries, k), dtype=np.uint64) + out_idx = np.zeros((n_queries, k), dtype=np.int64) out_dist = np.zeros((n_queries, k), dtype=np.float32) queries_device = device_ndarray(queries) @@ -358,7 +358,7 @@ def test_build_assertions(): index = ivf_flat.Index() queries = generate_data((n_queries, n_cols), np.float32) - out_idx = np.zeros((n_queries, k), dtype=np.uint64) + out_idx = np.zeros((n_queries, k), dtype=np.int64) out_dist = np.zeros((n_queries, k), dtype=np.float32) queries_device = device_ndarray(queries) @@ -381,7 +381,7 @@ def test_build_assertions(): index = ivf_flat.build(index_params, dataset_device) assert index.trained - indices = np.arange(n_rows + 1, dtype=np.uint64) + indices = np.arange(n_rows + 1, dtype=np.int64) indices_device = device_ndarray(indices) with pytest.raises(ValueError): @@ -424,7 +424,7 @@ def test_search_inputs(params): ).astype(q_dt, order=q_order) queries_device = device_ndarray(queries) - idx_dt = params.get("idx_dt", np.uint64) + idx_dt = params.get("idx_dt", np.int64) idx_order = params.get("idx_order", "C") out_idx = np.zeros( (params.get("idx_rows", n_queries), params.get("idx_cols", k)), From 5acdef1ac1ea01bf367076e98c6c741f5f7f92da Mon Sep 17 00:00:00 2001 From: divyegala Date: Tue, 14 Mar 2023 17:47:43 -0700 Subject: [PATCH 16/20] resolve bad merge, address review --- cpp/include/raft/neighbors/specializations.cuh | 5 ----- cpp/src/distance/neighbors/ivf_flat_build.cu | 2 +- cpp/src/distance/neighbors/ivf_flat_search.cu | 2 +- .../neighbors/specializations/ivfflat_build_float_int64_t.cu | 2 +- .../specializations/ivfflat_build_int8_t_int64_t.cu | 2 +- .../specializations/ivfflat_build_uint8_t_int64_t.cu | 2 +- .../specializations/ivfflat_extend_float_int64_t.cu | 2 +- .../specializations/ivfflat_extend_int8_t_int64_t.cu | 2 +- .../specializations/ivfflat_extend_uint8_t_int64_t.cu | 2 +- .../specializations/ivfflat_search_float_int64_t.cu | 2 +- .../specializations/ivfflat_search_int8_t_int64_t.cu | 2 +- .../specializations/ivfflat_search_uint8_t_int64_t.cu | 2 +- 12 files changed, 11 insertions(+), 16 deletions(-) diff --git a/cpp/include/raft/neighbors/specializations.cuh b/cpp/include/raft/neighbors/specializations.cuh index 42dc9a8087..27105b6eab 100644 --- a/cpp/include/raft/neighbors/specializations.cuh +++ b/cpp/include/raft/neighbors/specializations.cuh @@ -16,12 +16,7 @@ #pragma once -<<<<<<< HEAD -#include -#include #include -======= ->>>>>>> upstream/branch-23.04 #include #include diff --git a/cpp/src/distance/neighbors/ivf_flat_build.cu b/cpp/src/distance/neighbors/ivf_flat_build.cu index 9ceb455a9a..0d82fdbb08 100644 --- a/cpp/src/distance/neighbors/ivf_flat_build.cu +++ b/cpp/src/distance/neighbors/ivf_flat_build.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -#include +#include #include namespace raft::runtime::neighbors::ivf_flat { diff --git a/cpp/src/distance/neighbors/ivf_flat_search.cu b/cpp/src/distance/neighbors/ivf_flat_search.cu index f6d33df5be..b843ee7c30 100644 --- a/cpp/src/distance/neighbors/ivf_flat_search.cu +++ b/cpp/src/distance/neighbors/ivf_flat_search.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -#include +#include #include namespace raft::runtime::neighbors::ivf_flat { diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_build_float_int64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_build_float_int64_t.cu index a6c0723172..7082873d76 100644 --- a/cpp/src/distance/neighbors/specializations/ivfflat_build_float_int64_t.cu +++ b/cpp/src/distance/neighbors/specializations/ivfflat_build_float_int64_t.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -#include +#include namespace raft::neighbors::ivf_flat { diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_build_int8_t_int64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_build_int8_t_int64_t.cu index 50b7d069c9..ebc1a7fefa 100644 --- a/cpp/src/distance/neighbors/specializations/ivfflat_build_int8_t_int64_t.cu +++ b/cpp/src/distance/neighbors/specializations/ivfflat_build_int8_t_int64_t.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -#include +#include namespace raft::neighbors::ivf_flat { diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_build_uint8_t_int64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_build_uint8_t_int64_t.cu index d2ac83e770..870db6e97e 100644 --- a/cpp/src/distance/neighbors/specializations/ivfflat_build_uint8_t_int64_t.cu +++ b/cpp/src/distance/neighbors/specializations/ivfflat_build_uint8_t_int64_t.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -#include +#include namespace raft::neighbors::ivf_flat { diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_extend_float_int64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_extend_float_int64_t.cu index 5af3e37acd..71af06ad71 100644 --- a/cpp/src/distance/neighbors/specializations/ivfflat_extend_float_int64_t.cu +++ b/cpp/src/distance/neighbors/specializations/ivfflat_extend_float_int64_t.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -#include +#include namespace raft::neighbors::ivf_flat { diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_extend_int8_t_int64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_extend_int8_t_int64_t.cu index f89d523426..bb7bb6e7eb 100644 --- a/cpp/src/distance/neighbors/specializations/ivfflat_extend_int8_t_int64_t.cu +++ b/cpp/src/distance/neighbors/specializations/ivfflat_extend_int8_t_int64_t.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -#include +#include namespace raft::neighbors::ivf_flat { diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_extend_uint8_t_int64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_extend_uint8_t_int64_t.cu index 1fc8bca9ea..607b4b0913 100644 --- a/cpp/src/distance/neighbors/specializations/ivfflat_extend_uint8_t_int64_t.cu +++ b/cpp/src/distance/neighbors/specializations/ivfflat_extend_uint8_t_int64_t.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -#include +#include namespace raft::neighbors::ivf_flat { diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_search_float_int64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_search_float_int64_t.cu index 84d909d4bf..6de65546c8 100644 --- a/cpp/src/distance/neighbors/specializations/ivfflat_search_float_int64_t.cu +++ b/cpp/src/distance/neighbors/specializations/ivfflat_search_float_int64_t.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -#include +#include namespace raft::neighbors::ivf_flat { diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_search_int8_t_int64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_search_int8_t_int64_t.cu index 26cc91e224..8eda240ccd 100644 --- a/cpp/src/distance/neighbors/specializations/ivfflat_search_int8_t_int64_t.cu +++ b/cpp/src/distance/neighbors/specializations/ivfflat_search_int8_t_int64_t.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -#include +#include namespace raft::neighbors::ivf_flat { diff --git a/cpp/src/distance/neighbors/specializations/ivfflat_search_uint8_t_int64_t.cu b/cpp/src/distance/neighbors/specializations/ivfflat_search_uint8_t_int64_t.cu index 87de7183a0..8ff6533628 100644 --- a/cpp/src/distance/neighbors/specializations/ivfflat_search_uint8_t_int64_t.cu +++ b/cpp/src/distance/neighbors/specializations/ivfflat_search_uint8_t_int64_t.cu @@ -14,7 +14,7 @@ * limitations under the License. */ -#include +#include namespace raft::neighbors::ivf_flat { From 94611b273e28fc07ca98dec8d2a4554e89825108 Mon Sep 17 00:00:00 2001 From: divyegala Date: Tue, 14 Mar 2023 19:08:56 -0700 Subject: [PATCH 17/20] fix namespaces --- cpp/include/raft/neighbors/ivf_flat.cuh | 30 ++++++++++++------------- 1 file changed, 14 insertions(+), 16 deletions(-) diff --git a/cpp/include/raft/neighbors/ivf_flat.cuh b/cpp/include/raft/neighbors/ivf_flat.cuh index bc6a6d8b7b..30e0515e04 100644 --- a/cpp/include/raft/neighbors/ivf_flat.cuh +++ b/cpp/include/raft/neighbors/ivf_flat.cuh @@ -68,7 +68,7 @@ auto build(raft::device_resources const& handle, IdxT n_rows, uint32_t dim) -> index { - return raft::neighbors::ivf_flat::detail::build(handle, params, dataset, n_rows, dim); + return detail::build(handle, params, dataset, n_rows, dim); } /** @@ -113,11 +113,11 @@ auto build(raft::device_resources const& handle, raft::device_matrix_view dataset) -> index { - return raft::neighbors::ivf_flat::detail::build(handle, - params, - dataset.data_handle(), - static_cast(dataset.extent(0)), - static_cast(dataset.extent(1))); + return detail::build(handle, + params, + dataset.data_handle(), + static_cast(dataset.extent(0)), + static_cast(dataset.extent(1))); } /** @@ -159,11 +159,11 @@ void build(raft::device_resources const& handle, raft::device_matrix_view dataset, raft::neighbors::ivf_flat::index& idx) { - idx = raft::spatial::knn::ivf_flat::detail::build(handle, - params, - dataset.data_handle(), - static_cast(dataset.extent(0)), - static_cast(dataset.extent(1))); + idx = detail::build(handle, + params, + dataset.data_handle(), + static_cast(dataset.extent(0)), + static_cast(dataset.extent(1))); } /** @} */ @@ -207,8 +207,7 @@ auto extend(raft::device_resources const& handle, const IdxT* new_indices, IdxT n_rows) -> index { - return raft::neighbors::ivf_flat::detail::extend( - handle, orig_index, new_vectors, new_indices, n_rows); + return detail::extend(handle, orig_index, new_vectors, new_indices, n_rows); } /** @@ -296,7 +295,7 @@ void extend(raft::device_resources const& handle, const IdxT* new_indices, IdxT n_rows) { - raft::neighbors::ivf_flat::detail::extend(handle, index, new_vectors, new_indices, n_rows); + detail::extend(handle, index, new_vectors, new_indices, n_rows); } /** @@ -398,8 +397,7 @@ void search(raft::device_resources const& handle, float* distances, rmm::mr::device_memory_resource* mr = nullptr) { - return raft::neighbors::ivf_flat::detail::search( - handle, params, index, queries, n_queries, k, neighbors, distances, mr); + return detail::search(handle, params, index, queries, n_queries, k, neighbors, distances, mr); } /** From a8a94412b55a333e12eb1a1dc3e4a6b7a0112f04 Mon Sep 17 00:00:00 2001 From: divyegala Date: Wed, 15 Mar 2023 07:13:15 -0700 Subject: [PATCH 18/20] add missing TU back --- cpp/CMakeLists.txt | 3 ++- cpp/include/raft/neighbors/ivf_flat.cuh | 30 +++++++++++++------------ 2 files changed, 18 insertions(+), 15 deletions(-) diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index 64739b5dda..76bc551b8b 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -362,6 +362,8 @@ if(RAFT_COMPILE_DIST_LIBRARY) src/distance/distance/specializations/fused_l2_nn_float_int64.cu src/distance/matrix/specializations/detail/select_k_float_uint32_t.cu src/distance/matrix/specializations/detail/select_k_float_int64_t.cu + src/distance/matrix/specializations/detail/select_k_half_uint32_t.cu + src/distance/matrix/specializations/detail/select_k_half_int64_t.cu src/distance/neighbors/ivf_flat_search.cu src/distance/neighbors/ivf_flat_build.cu src/distance/neighbors/specializations/ivfflat_build_float_int64_t.cu @@ -373,7 +375,6 @@ if(RAFT_COMPILE_DIST_LIBRARY) src/distance/neighbors/specializations/ivfflat_search_float_int64_t.cu src/distance/neighbors/specializations/ivfflat_search_int8_t_int64_t.cu src/distance/neighbors/specializations/ivfflat_search_uint8_t_int64_t.cu - src/distance/matrix/specializations/detail/select_k_half_int64_t.cu src/distance/neighbors/ivfpq_build.cu src/distance/neighbors/ivfpq_deserialize.cu src/distance/neighbors/ivfpq_serialize.cu diff --git a/cpp/include/raft/neighbors/ivf_flat.cuh b/cpp/include/raft/neighbors/ivf_flat.cuh index 30e0515e04..c573676504 100644 --- a/cpp/include/raft/neighbors/ivf_flat.cuh +++ b/cpp/include/raft/neighbors/ivf_flat.cuh @@ -68,7 +68,7 @@ auto build(raft::device_resources const& handle, IdxT n_rows, uint32_t dim) -> index { - return detail::build(handle, params, dataset, n_rows, dim); + return raft::neighbors::ivf_flat::detail::build(handle, params, dataset, n_rows, dim); } /** @@ -113,11 +113,11 @@ auto build(raft::device_resources const& handle, raft::device_matrix_view dataset) -> index { - return detail::build(handle, - params, - dataset.data_handle(), - static_cast(dataset.extent(0)), - static_cast(dataset.extent(1))); + return raft::neighbors::ivf_flat::detail::build(handle, + params, + dataset.data_handle(), + static_cast(dataset.extent(0)), + static_cast(dataset.extent(1))); } /** @@ -159,11 +159,11 @@ void build(raft::device_resources const& handle, raft::device_matrix_view dataset, raft::neighbors::ivf_flat::index& idx) { - idx = detail::build(handle, - params, - dataset.data_handle(), - static_cast(dataset.extent(0)), - static_cast(dataset.extent(1))); + idx = raft::neighbors::ivf_flat::detail::build(handle, + params, + dataset.data_handle(), + static_cast(dataset.extent(0)), + static_cast(dataset.extent(1))); } /** @} */ @@ -207,7 +207,8 @@ auto extend(raft::device_resources const& handle, const IdxT* new_indices, IdxT n_rows) -> index { - return detail::extend(handle, orig_index, new_vectors, new_indices, n_rows); + return raft::neighbors::ivf_flat::detail::extend( + handle, orig_index, new_vectors, new_indices, n_rows); } /** @@ -295,7 +296,7 @@ void extend(raft::device_resources const& handle, const IdxT* new_indices, IdxT n_rows) { - detail::extend(handle, index, new_vectors, new_indices, n_rows); + raft::neighbors::ivf_flat::detail::extend(handle, index, new_vectors, new_indices, n_rows); } /** @@ -397,7 +398,8 @@ void search(raft::device_resources const& handle, float* distances, rmm::mr::device_memory_resource* mr = nullptr) { - return detail::search(handle, params, index, queries, n_queries, k, neighbors, distances, mr); + return raft::neighbors::ivf_flat::detail::search( + handle, params, index, queries, n_queries, k, neighbors, distances, mr); } /** From d1d1c790d8fcc55527ae139f92e4e70da7b19bff Mon Sep 17 00:00:00 2001 From: divyegala Date: Wed, 15 Mar 2023 11:12:56 -0700 Subject: [PATCH 19/20] fix index --- cpp/include/raft/neighbors/ivf_flat_types.hpp | 8 ++++---- .../neighbors/ivf_flat/cpp/c_ivf_flat.pxd | 5 +++-- .../pylibraft/neighbors/ivf_flat/ivf_flat.pyx | 15 +++++++++------ 3 files changed, 16 insertions(+), 12 deletions(-) diff --git a/cpp/include/raft/neighbors/ivf_flat_types.hpp b/cpp/include/raft/neighbors/ivf_flat_types.hpp index 20bca6e3e6..db9cf3b6f1 100644 --- a/cpp/include/raft/neighbors/ivf_flat_types.hpp +++ b/cpp/include/raft/neighbors/ivf_flat_types.hpp @@ -237,9 +237,9 @@ struct index : ann::index { index(raft::device_resources const& res, raft::distance::DistanceType metric, uint32_t n_lists, + uint32_t dim, bool adaptive_centers, - bool conservative_memory_allocation, - uint32_t dim) + bool conservative_memory_allocation) : ann::index(), veclen_(calculate_veclen(dim)), metric_(metric), @@ -261,9 +261,9 @@ struct index : ann::index { : index(res, params.metric, params.n_lists, + dim, params.adaptive_centers, - params.conservative_memory_allocation, - dim) + params.conservative_memory_allocation) { } diff --git a/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/c_ivf_flat.pxd b/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/c_ivf_flat.pxd index 94f7cfceb6..6e56ca39d7 100644 --- a/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/c_ivf_flat.pxd +++ b/python/pylibraft/pylibraft/neighbors/ivf_flat/cpp/c_ivf_flat.pxd @@ -54,15 +54,16 @@ cdef extern from "raft/neighbors/ivf_flat_types.hpp" \ uint32_t n_lists uint32_t kmeans_n_iters double kmeans_trainset_fraction - bool add_data_on_build bool adaptive_centers + bool conservative_memory_allocation cdef cppclass index[T, IdxT](ann_index): index(const device_resources& handle, DistanceType metric, uint32_t n_lists, + uint32_t dim, bool adaptive_centers, - uint32_t dim) + bool conservative_memory_allocation) IdxT size() uint32_t dim() DistanceType metric() diff --git a/python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx b/python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx index b89ac3ada5..988599f35b 100644 --- a/python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx +++ b/python/pylibraft/pylibraft/neighbors/ivf_flat/ivf_flat.pyx @@ -180,8 +180,9 @@ cdef class IndexFloat(Index): self.index = new c_ivf_flat.index[float, int64_t]( deref(handle_), _get_metric("sqeuclidean"), 1, - False, - 8) + 4, + False, + False) def __repr__(self): m_str = "metric=" + _get_metric_string(self.index.metric()) @@ -228,8 +229,9 @@ cdef class IndexInt8(Index): self.index = new c_ivf_flat.index[int8_t, int64_t]( deref(handle_), _get_metric("sqeuclidean"), 1, - False, - 8) + 4, + False, + False) def __repr__(self): m_str = "metric=" + _get_metric_string(self.index.metric()) @@ -276,8 +278,9 @@ cdef class IndexUint8(Index): self.index = new c_ivf_flat.index[uint8_t, int64_t]( deref(handle_), _get_metric("sqeuclidean"), 1, - False, - 8) + 4, + False, + False) def __repr__(self): m_str = "metric=" + _get_metric_string(self.index.metric()) From da39e91236e6b303487b48e741b812c8870ade81 Mon Sep 17 00:00:00 2001 From: divyegala Date: Thu, 16 Mar 2023 14:58:37 -0700 Subject: [PATCH 20/20] all tests passing --- cpp/test/neighbors/ann_ivf_flat.cuh | 2 +- cpp/test/neighbors/refine.cu | 2 +- python/pylibraft/setup.py | 15 +-------------- 3 files changed, 3 insertions(+), 16 deletions(-) diff --git a/cpp/test/neighbors/ann_ivf_flat.cuh b/cpp/test/neighbors/ann_ivf_flat.cuh index ac8c37542c..486ff61724 100644 --- a/cpp/test/neighbors/ann_ivf_flat.cuh +++ b/cpp/test/neighbors/ann_ivf_flat.cuh @@ -192,7 +192,7 @@ class AnnIVFFlatTest : public ::testing::TestWithParam> { new_half_of_data_view, std::make_optional>( new_half_of_data_indices_view), - index_2); + &index_2); auto search_queries_view = raft::make_device_matrix_view( search_queries.data(), ps.num_queries, ps.dim); diff --git a/cpp/test/neighbors/refine.cu b/cpp/test/neighbors/refine.cu index 5b30045928..8866c404a9 100644 --- a/cpp/test/neighbors/refine.cu +++ b/cpp/test/neighbors/refine.cu @@ -115,7 +115,7 @@ const std::vector> inputs = {static_cast(1), static_cast(10), static_cast(33)}, {static_cast(33)}, {raft::distance::DistanceType::L2Expanded, raft::distance::DistanceType::InnerProduct}, - {true}); + {false, true}); typedef RefineTest RefineTestF; TEST_P(RefineTestF, AnnRefine) { this->testRefine(); } diff --git a/python/pylibraft/setup.py b/python/pylibraft/setup.py index c2e7c29631..d00e36d804 100644 --- a/python/pylibraft/setup.py +++ b/python/pylibraft/setup.py @@ -28,22 +28,9 @@ def exclude_libcxx_symlink(cmake_manifest): setup( - # include_package_data=True, + include_package_data=True, # Don't want libcxx getting pulled into wheel builds. cmake_process_manifest_hook=exclude_libcxx_symlink, packages=find_packages(include=["pylibraft", "pylibraft.*"]), zip_safe=False, - package_data={ - key: ["*.hpp", "*.pxd"] - for key in find_packages( - include=[ - "pylibraft.distance", - "pylibraft.distance.includes", - "pylibraft.common", - "pylibraft.common.includes", - "pylibraft.random", - "pylibraft.random.includes", - ] - ) - }, )