From f9e5165d3019f43dc959167ba21a2ac0827d4403 Mon Sep 17 00:00:00 2001 From: "Corey J. Nolet" Date: Wed, 26 Oct 2022 15:51:35 -0400 Subject: [PATCH 1/6] iFixing ivf flat test so it actually invokes mdspan apis --- cpp/include/raft/neighbors/ivf_flat.cuh | 22 +++++++------- cpp/include/raft/spatial/knn/ivf_flat.cuh | 2 +- .../raft/spatial/knn/ivf_flat_types.hpp | 2 +- cpp/test/neighbors/ann_ivf_flat.cu | 29 +++++++++++++------ 4 files changed, 33 insertions(+), 22 deletions(-) diff --git a/cpp/include/raft/neighbors/ivf_flat.cuh b/cpp/include/raft/neighbors/ivf_flat.cuh index 87400a9b93..677c16bc01 100644 --- a/cpp/include/raft/neighbors/ivf_flat.cuh +++ b/cpp/include/raft/neighbors/ivf_flat.cuh @@ -195,7 +195,7 @@ auto extend(const handle_t& handle, std::optional> new_indices = std::nullopt) -> index { - return raft::spatial::knn::ivf_flat::detail::extend( + return raft::spatial::knn::ivf_flat::extend( handle, orig_index, new_vectors.data_handle(), @@ -376,7 +376,7 @@ void search(const handle_t& handle, const index& index, raft::device_matrix_view queries, raft::device_matrix_view neighbors, - raft::device_matrix_view distances, + raft::device_matrix_view distances, const search_params& params, int_t k) { @@ -391,15 +391,15 @@ void search(const handle_t& handle, RAFT_EXPECTS(queries.extent(1) == index.dim(), "Number of query dimensions should equal number of dimensions in the index."); - return raft::spatial::knn::ivf_flat::detail::search(handle, - params, - index, - queries.data_handle(), - queries.extent(0), - k, - neighbors.data_handle(), - distances.data_handle(), - nullptr); + return raft::spatial::knn::ivf_flat::search(handle, + params, + index, + queries.data_handle(), + const_cast(queries.extent(0)), + const_cast(k), + neighbors.data_handle(), + distances.data_handle(), + nullptr); } } // namespace raft::neighbors::ivf_flat diff --git a/cpp/include/raft/spatial/knn/ivf_flat.cuh b/cpp/include/raft/spatial/knn/ivf_flat.cuh index 65b6f5ed4b..92fe49be98 100644 --- a/cpp/include/raft/spatial/knn/ivf_flat.cuh +++ b/cpp/include/raft/spatial/knn/ivf_flat.cuh @@ -36,4 +36,4 @@ using raft::neighbors::ivf_flat::build; using raft::neighbors::ivf_flat::extend; using raft::neighbors::ivf_flat::search; -} // namespace raft::spatial::knn::ivf_flat +}; // namespace raft::spatial::knn::ivf_flat diff --git a/cpp/include/raft/spatial/knn/ivf_flat_types.hpp b/cpp/include/raft/spatial/knn/ivf_flat_types.hpp index 2db29eeb58..75d777573f 100644 --- a/cpp/include/raft/spatial/knn/ivf_flat_types.hpp +++ b/cpp/include/raft/spatial/knn/ivf_flat_types.hpp @@ -37,4 +37,4 @@ using raft::neighbors::ivf_flat::index_params; using raft::neighbors::ivf_flat::kIndexGroupSize; using raft::neighbors::ivf_flat::search_params; -} // namespace raft::spatial::knn::ivf_flat +}; // namespace raft::spatial::knn::ivf_flat diff --git a/cpp/test/neighbors/ann_ivf_flat.cu b/cpp/test/neighbors/ann_ivf_flat.cu index 3a5daff4bb..d32b0e56f6 100644 --- a/cpp/test/neighbors/ann_ivf_flat.cu +++ b/cpp/test/neighbors/ann_ivf_flat.cu @@ -169,20 +169,31 @@ class AnnIVFFlatTest : public ::testing::TestWithParam> { auto index_2 = ivf_flat::extend(handle_, index, half_of_data_view); + + 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 + ); + + auto new_half_of_data_indices_view = raft::make_device_vector_view( + vector_indices.data() + half_of_data, IdxT(ps.num_db_vecs) - half_of_data + ); + + ivf_flat::extend(handle_, &index_2, - database.data() + half_of_data * ps.dim, - vector_indices.data() + half_of_data, - IdxT(ps.num_db_vecs) - half_of_data); + new_half_of_data_view, + std::make_optional>(new_half_of_data_indices_view)); + auto search_queries_view = raft::make_device_matrix_view(search_queries.data(), ps.num_queries, ps.dim); + auto indices_out_view = raft::make_device_matrix_view(indices_ivfflat_dev.data(), ps.num_queries, ps.k); + auto dists_out_view = raft::make_device_matrix_view(distances_ivfflat_dev.data(), ps.num_queries, ps.k); ivf_flat::search(handle_, - search_params, index_2, - search_queries.data(), - ps.num_queries, - ps.k, - indices_ivfflat_dev.data(), - distances_ivfflat_dev.data()); + search_queries_view, + indices_out_view, + dists_out_view, + search_params, + ps.k); update_host(distances_ivfflat.data(), distances_ivfflat_dev.data(), queries_size, stream_); update_host(indices_ivfflat.data(), indices_ivfflat_dev.data(), queries_size, stream_); From 7ebe72c489a54b1b771a704151bfc969fc0ce7a0 Mon Sep 17 00:00:00 2001 From: "Corey J. Nolet" Date: Wed, 26 Oct 2022 16:18:19 -0400 Subject: [PATCH 2/6] Fixing namespace --- cpp/include/raft/neighbors/ivf_flat.cuh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/include/raft/neighbors/ivf_flat.cuh b/cpp/include/raft/neighbors/ivf_flat.cuh index 677c16bc01..c61b19fd01 100644 --- a/cpp/include/raft/neighbors/ivf_flat.cuh +++ b/cpp/include/raft/neighbors/ivf_flat.cuh @@ -195,7 +195,7 @@ auto extend(const handle_t& handle, std::optional> new_indices = std::nullopt) -> index { - return raft::spatial::knn::ivf_flat::extend( + return extend( handle, orig_index, new_vectors.data_handle(), @@ -391,7 +391,7 @@ void search(const handle_t& handle, RAFT_EXPECTS(queries.extent(1) == index.dim(), "Number of query dimensions should equal number of dimensions in the index."); - return raft::spatial::knn::ivf_flat::search(handle, + return search(handle, params, index, queries.data_handle(), From d8a926be2c7a9724c12374bf7070fccb0f08fd78 Mon Sep 17 00:00:00 2001 From: "Corey J. Nolet" Date: Wed, 26 Oct 2022 16:19:29 -0400 Subject: [PATCH 3/6] Fixing cast --- cpp/include/raft/neighbors/ivf_flat.cuh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/include/raft/neighbors/ivf_flat.cuh b/cpp/include/raft/neighbors/ivf_flat.cuh index c61b19fd01..0c109f6631 100644 --- a/cpp/include/raft/neighbors/ivf_flat.cuh +++ b/cpp/include/raft/neighbors/ivf_flat.cuh @@ -395,8 +395,8 @@ void search(const handle_t& handle, params, index, queries.data_handle(), - const_cast(queries.extent(0)), - const_cast(k), + static(queries.extent(0)), + static(k), neighbors.data_handle(), distances.data_handle(), nullptr); From cacb6f3dd5c38238b7099f74daaa0871c7502866 Mon Sep 17 00:00:00 2001 From: "Corey J. Nolet" Date: Wed, 26 Oct 2022 16:20:16 -0400 Subject: [PATCH 4/6] Correct casting --- cpp/include/raft/neighbors/ivf_flat.cuh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/include/raft/neighbors/ivf_flat.cuh b/cpp/include/raft/neighbors/ivf_flat.cuh index 0c109f6631..e9b38582ab 100644 --- a/cpp/include/raft/neighbors/ivf_flat.cuh +++ b/cpp/include/raft/neighbors/ivf_flat.cuh @@ -395,8 +395,8 @@ void search(const handle_t& handle, params, index, queries.data_handle(), - static(queries.extent(0)), - static(k), + static_cast(queries.extent(0)), + static_cast(k), neighbors.data_handle(), distances.data_handle(), nullptr); From db9f965c662f8187d5c18e24ae5fbbea0715ea33 Mon Sep 17 00:00:00 2001 From: "Corey J. Nolet" Date: Wed, 26 Oct 2022 20:48:58 -0400 Subject: [PATCH 5/6] Important updates --- cpp/include/raft/neighbors/ivf_flat.cuh | 16 ++++++++-------- cpp/test/neighbors/ann_ivf_flat.cu | 22 +++++++++++----------- 2 files changed, 19 insertions(+), 19 deletions(-) diff --git a/cpp/include/raft/neighbors/ivf_flat.cuh b/cpp/include/raft/neighbors/ivf_flat.cuh index e9b38582ab..0f8b7a9b6b 100644 --- a/cpp/include/raft/neighbors/ivf_flat.cuh +++ b/cpp/include/raft/neighbors/ivf_flat.cuh @@ -392,14 +392,14 @@ void search(const handle_t& handle, "Number of query dimensions should equal number of dimensions in the index."); return search(handle, - params, - index, - queries.data_handle(), - static_cast(queries.extent(0)), - static_cast(k), - neighbors.data_handle(), - distances.data_handle(), - nullptr); + params, + index, + queries.data_handle(), + static_cast(queries.extent(0)), + static_cast(k), + neighbors.data_handle(), + distances.data_handle(), + nullptr); } } // namespace raft::neighbors::ivf_flat diff --git a/cpp/test/neighbors/ann_ivf_flat.cu b/cpp/test/neighbors/ann_ivf_flat.cu index d32b0e56f6..9a430e14f2 100644 --- a/cpp/test/neighbors/ann_ivf_flat.cu +++ b/cpp/test/neighbors/ann_ivf_flat.cu @@ -169,24 +169,24 @@ class AnnIVFFlatTest : public ::testing::TestWithParam> { auto index_2 = ivf_flat::extend(handle_, index, half_of_data_view); - 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 - ); + database.data() + half_of_data * ps.dim, IdxT(ps.num_db_vecs) - half_of_data, ps.dim); auto new_half_of_data_indices_view = raft::make_device_vector_view( - vector_indices.data() + half_of_data, IdxT(ps.num_db_vecs) - half_of_data - ); - + 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)); - - auto search_queries_view = raft::make_device_matrix_view(search_queries.data(), ps.num_queries, ps.dim); - auto indices_out_view = raft::make_device_matrix_view(indices_ivfflat_dev.data(), ps.num_queries, ps.k); - auto dists_out_view = raft::make_device_matrix_view(distances_ivfflat_dev.data(), ps.num_queries, ps.k); + std::make_optional>( + new_half_of_data_indices_view)); + + auto search_queries_view = raft::make_device_matrix_view( + search_queries.data(), ps.num_queries, ps.dim); + auto indices_out_view = raft::make_device_matrix_view( + indices_ivfflat_dev.data(), ps.num_queries, ps.k); + auto dists_out_view = raft::make_device_matrix_view( + distances_ivfflat_dev.data(), ps.num_queries, ps.k); ivf_flat::search(handle_, index_2, search_queries_view, From 0f96bbdec04165607b0a0e1743658efbadb40c0a Mon Sep 17 00:00:00 2001 From: "Corey J. Nolet" Date: Wed, 26 Oct 2022 20:51:52 -0400 Subject: [PATCH 6/6] Updating docs --- cpp/include/raft/neighbors/ivf_flat.cuh | 8 ++++---- cpp/include/raft/neighbors/ivf_pq.cuh | 2 +- 2 files changed, 5 insertions(+), 5 deletions(-) diff --git a/cpp/include/raft/neighbors/ivf_flat.cuh b/cpp/include/raft/neighbors/ivf_flat.cuh index 0f8b7a9b6b..5317f406e1 100644 --- a/cpp/include/raft/neighbors/ivf_flat.cuh +++ b/cpp/include/raft/neighbors/ivf_flat.cuh @@ -121,7 +121,7 @@ auto build(const handle_t& handle, * * Usage example: * @code{.cpp} - * using namespace raft::spatial::knn; + * using namespace raft::neighbors; * ivf_flat::index_params index_params; * index_params.add_data_on_build = false; // don't populate index on build * index_params.kmeans_trainset_fraction = 1.0; // use whole dataset for kmeans training @@ -164,7 +164,7 @@ auto extend(const handle_t& handle, * * Usage example: * @code{.cpp} - * using namespace raft::spatial::knn; + * using namespace raft::neighbors; * ivf_flat::index_params index_params; * index_params.add_data_on_build = false; // don't populate index on build * index_params.kmeans_trainset_fraction = 1.0; // use whole dataset for kmeans training @@ -208,7 +208,7 @@ auto extend(const handle_t& handle, * * Usage example: * @code{.cpp} - * using namespace raft::spatial::knn; + * using namespace raft::neighbors; * ivf_flat::index_params index_params; * index_params.add_data_on_build = false; // don't populate index on build * index_params.kmeans_trainset_fraction = 1.0; // use whole dataset for kmeans training @@ -244,7 +244,7 @@ void extend(const handle_t& handle, * * Usage example: * @code{.cpp} - * using namespace raft::spatial::knn; + * using namespace raft::neighbors; * ivf_flat::index_params index_params; * index_params.add_data_on_build = false; // don't populate index on build * index_params.kmeans_trainset_fraction = 1.0; // use whole dataset for kmeans training diff --git a/cpp/include/raft/neighbors/ivf_pq.cuh b/cpp/include/raft/neighbors/ivf_pq.cuh index 5d619c5bec..207e298947 100644 --- a/cpp/include/raft/neighbors/ivf_pq.cuh +++ b/cpp/include/raft/neighbors/ivf_pq.cuh @@ -76,7 +76,7 @@ inline auto build( * * Usage example: * @code{.cpp} - * using namespace raft::spatial::knn; + * using namespace raft::neighbors; * ivf_pq::index_params index_params; * index_params.add_data_on_build = false; // don't populate index on build * index_params.kmeans_trainset_fraction = 1.0; // use whole dataset for kmeans training