From ba244fb9db1faabcb9e0875d583882ecba4a280f Mon Sep 17 00:00:00 2001 From: rhdong Date: Sun, 26 Jan 2025 22:58:37 -0800 Subject: [PATCH 01/10] [Feat] Add Support for Index `merge` in CAGRA --- cpp/CMakeLists.txt | 4 + cpp/include/cuvs/neighbors/cagra.hpp | 146 ++++++++++++++ cpp/src/neighbors/cagra.cuh | 9 + cpp/src/neighbors/cagra_merge_float.cu | 35 ++++ cpp/src/neighbors/cagra_merge_half.cu | 35 ++++ cpp/src/neighbors/cagra_merge_int8.cu | 35 ++++ cpp/src/neighbors/cagra_merge_uint8.cu | 35 ++++ .../neighbors/detail/cagra/cagra_merge.cuh | 130 +++++++++++++ cpp/test/neighbors/ann_cagra.cuh | 184 ++++++++++++++++++ .../ann_cagra/test_float_uint32_t.cu | 6 + .../neighbors/ann_cagra/test_half_uint32_t.cu | 6 + .../ann_cagra/test_int8_t_uint32_t.cu | 6 +- .../ann_cagra/test_uint8_t_uint32_t.cu | 5 + 13 files changed, 635 insertions(+), 1 deletion(-) create mode 100644 cpp/src/neighbors/cagra_merge_float.cu create mode 100644 cpp/src/neighbors/cagra_merge_half.cu create mode 100644 cpp/src/neighbors/cagra_merge_int8.cu create mode 100644 cpp/src/neighbors/cagra_merge_uint8.cu create mode 100644 cpp/src/neighbors/detail/cagra/cagra_merge.cuh diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index f88ca1af51..7b6f9e3e27 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -390,6 +390,10 @@ if(BUILD_SHARED_LIBS) src/neighbors/cagra_serialize_half.cu src/neighbors/cagra_serialize_int8.cu src/neighbors/cagra_serialize_uint8.cu + src/neighbors/cagra_merge_float.cu + src/neighbors/cagra_merge_half.cu + src/neighbors/cagra_merge_int8.cu + src/neighbors/cagra_merge_uint8.cu src/neighbors/iface/iface_cagra_float_uint32_t.cu src/neighbors/iface/iface_cagra_half_uint32_t.cu src/neighbors/iface/iface_cagra_int8_t_uint32_t.cu diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index a4684ce267..8f257900ea 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -248,6 +248,17 @@ struct extend_params { * 0. */ uint32_t max_chunk_size = 0; }; +/** + * @} + */ + +/** + * @defgroup cagra_cpp_merge_params CAGRA index merge parameters + * @{ + */ +struct merge_params : public index_params { + merge_params(const index_params& params) : index_params(params) {} +}; /** * @} @@ -1747,7 +1758,142 @@ void serialize_to_hnswlib(raft::resources const& handle, void serialize_to_hnswlib(raft::resources const& handle, const std::string& filename, const cuvs::neighbors::cagra::index& index); +/** + * @} + */ + +/** + * @defgroup cagra_cpp_index_merge CAGRA index build functions + * @{ + */ + +/** @brief Merge multiple CAGRA indices into a single index. + * + * This function merges multiple CAGRA indices into one, combining both the datasets and graph + * structures. + * + * Usage example: + * @code{.cpp} + * using namespace raft::neighbors; + * auto dataset0 = raft::make_host_matrix(handle, size0, dim); + * auto dataset1 = raft::make_host_matrix(handle, size1, dim); + * + * auto index0 = cagra::build(res, index_params, dataset0); + * auto index1 = cagra::build(res, index_params, dataset1); + * + * std::vector*> indices{&index0, &index1}; + * cagra::merge_params params{index_params}; + * + * auto merged_index = cagra::merge(res, params, indices); + * @endcode + * + * @param[in] res RAFT resources used for the merge operation. + * @param[in] params Parameters that control the merging process. + * @param[in] indices A vector of pointers to the CAGRA indices to merge. All indices must: + * - Have attached datasets with the same dimensionality. + * + * @return A new CAGRA index containing the merged indices, graph, and dataset. + */ +auto merge(raft::resources const& res, + const cuvs::neighbors::cagra::merge_params& params, + std::vector*>& indices) + -> cuvs::neighbors::cagra::index; + +/** @brief Merge multiple CAGRA indices into a single index. + * + * This function merges multiple CAGRA indices into one, combining both the datasets and graph + * structures. + * + * Usage example: + * @code{.cpp} + * using namespace raft::neighbors; + * auto dataset0 = raft::make_host_matrix(handle, size0, dim); + * auto dataset1 = raft::make_host_matrix(handle, size1, dim); + * + * auto index0 = cagra::build(res, index_params, dataset0); + * auto index1 = cagra::build(res, index_params, dataset1); + * + * std::vector*> indices{&index0, &index1}; + * cagra::merge_params params{index_params}; + * + * auto merged_index = cagra::merge(res, params, indices); + * @endcode + * + * @param[in] res RAFT resources used for the merge operation. + * @param[in] params Parameters that control the merging process. + * @param[in] indices A vector of pointers to the CAGRA indices to merge. All indices must: + * - Have attached datasets with the same dimensionality. + * + * @return A new CAGRA index containing the merged indices, graph, and dataset. + */ +auto merge(raft::resources const& res, + const cuvs::neighbors::cagra::merge_params& params, + std::vector*>& indices) + -> cuvs::neighbors::cagra::index; +/** @brief Merge multiple CAGRA indices into a single index. + * + * This function merges multiple CAGRA indices into one, combining both the datasets and graph + * structures. + * + * Usage example: + * @code{.cpp} + * using namespace raft::neighbors; + * auto dataset0 = raft::make_host_matrix(handle, size0, dim); + * auto dataset1 = raft::make_host_matrix(handle, size1, dim); + * + * auto index0 = cagra::build(res, index_params, dataset0); + * auto index1 = cagra::build(res, index_params, dataset1); + * + * std::vector*> indices{&index0, &index1}; + * cagra::merge_params params{index_params}; + * + * auto merged_index = cagra::merge(res, params, indices); + * @endcode + * + * @param[in] res RAFT resources used for the merge operation. + * @param[in] params Parameters that control the merging process. + * @param[in] indices A vector of pointers to the CAGRA indices to merge. All indices must: + * - Have attached datasets with the same dimensionality. + * + * @return A new CAGRA index containing the merged indices, graph, and dataset. + */ +auto merge(raft::resources const& res, + const cuvs::neighbors::cagra::merge_params& params, + std::vector*>& indices) + -> cuvs::neighbors::cagra::index; + +/** @brief Merge multiple CAGRA indices into a single index. + * + * This function merges multiple CAGRA indices into one, combining both the datasets and graph + * structures. + * + * Usage example: + * @code{.cpp} + * using namespace raft::neighbors; + * auto dataset0 = raft::make_host_matrix(handle, size0, dim); + * auto dataset1 = raft::make_host_matrix(handle, size1, dim); + * + * auto index0 = cagra::build(res, index_params, dataset0); + * auto index1 = cagra::build(res, index_params, dataset1); + * + * std::vector*> indices{&index0, &index1}; + * cagra::merge_params params{index_params}; + * + * auto merged_index = cagra::merge(res, params, indices); + * @endcode + * + * @param[in] res RAFT resources used for the merge operation. + * @param[in] params Parameters that control the merging process. + * @param[in] indices A vector of pointers to the CAGRA indices to merge. All indices must: + * - Have attached datasets with the same dimensionality. + * + * @return A new CAGRA index containing the merged indices, graph, and dataset. + */ +auto merge(raft::resources const& res, + const cuvs::neighbors::cagra::merge_params& params, + std::vector*>& indices) + -> cuvs::neighbors::cagra::index; /** * @} */ diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index dacfd6f630..0ec419ec2c 100644 --- a/cpp/src/neighbors/cagra.cuh +++ b/cpp/src/neighbors/cagra.cuh @@ -18,6 +18,7 @@ #include "detail/cagra/add_nodes.cuh" #include "detail/cagra/cagra_build.cuh" +#include "detail/cagra/cagra_merge.cuh" #include "detail/cagra/cagra_search.cuh" #include "detail/cagra/graph_core.cuh" @@ -369,6 +370,14 @@ void extend( cagra::extend_core(handle, additional_dataset, index, params, ndv, ngv); } +template +index merge(raft::resources const& handle, + const cagra::merge_params& params, + std::vector*>& indices) +{ + return cagra::detail::merge(handle, params, indices); +} + /** @} */ // end group cagra } // namespace cuvs::neighbors::cagra diff --git a/cpp/src/neighbors/cagra_merge_float.cu b/cpp/src/neighbors/cagra_merge_float.cu new file mode 100644 index 0000000000..951c0c5fe9 --- /dev/null +++ b/cpp/src/neighbors/cagra_merge_float.cu @@ -0,0 +1,35 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "cagra.cuh" +#include + +namespace cuvs::neighbors::cagra { + +#define RAFT_INST_CAGRA_MERGE(T, IdxT) \ + auto merge(raft::resources const& handle, \ + const cuvs::neighbors::cagra::merge_params& params, \ + std::vector*>& indices) \ + ->cuvs::neighbors::cagra::index \ + { \ + return cuvs::neighbors::cagra::merge(handle, params, indices); \ + } + +RAFT_INST_CAGRA_MERGE(float, uint32_t); + +#undef RAFT_INST_CAGRA_MERGE + +} // namespace cuvs::neighbors::cagra diff --git a/cpp/src/neighbors/cagra_merge_half.cu b/cpp/src/neighbors/cagra_merge_half.cu new file mode 100644 index 0000000000..704a00f747 --- /dev/null +++ b/cpp/src/neighbors/cagra_merge_half.cu @@ -0,0 +1,35 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "cagra.cuh" +#include + +namespace cuvs::neighbors::cagra { + +#define RAFT_INST_CAGRA_MERGE(T, IdxT) \ + auto merge(raft::resources const& handle, \ + const cuvs::neighbors::cagra::merge_params& params, \ + std::vector*>& indices) \ + ->cuvs::neighbors::cagra::index \ + { \ + return cuvs::neighbors::cagra::merge(handle, params, indices); \ + } + +RAFT_INST_CAGRA_MERGE(half, uint32_t); + +#undef RAFT_INST_CAGRA_MERGE + +} // namespace cuvs::neighbors::cagra diff --git a/cpp/src/neighbors/cagra_merge_int8.cu b/cpp/src/neighbors/cagra_merge_int8.cu new file mode 100644 index 0000000000..a7e9035626 --- /dev/null +++ b/cpp/src/neighbors/cagra_merge_int8.cu @@ -0,0 +1,35 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "cagra.cuh" +#include + +namespace cuvs::neighbors::cagra { + +#define RAFT_INST_CAGRA_MERGE(T, IdxT) \ + auto merge(raft::resources const& handle, \ + const cuvs::neighbors::cagra::merge_params& params, \ + std::vector*>& indices) \ + ->cuvs::neighbors::cagra::index \ + { \ + return cuvs::neighbors::cagra::merge(handle, params, indices); \ + } + +RAFT_INST_CAGRA_MERGE(int8_t, uint32_t); + +#undef RAFT_INST_CAGRA_MERGE + +} // namespace cuvs::neighbors::cagra diff --git a/cpp/src/neighbors/cagra_merge_uint8.cu b/cpp/src/neighbors/cagra_merge_uint8.cu new file mode 100644 index 0000000000..a4fc7149c1 --- /dev/null +++ b/cpp/src/neighbors/cagra_merge_uint8.cu @@ -0,0 +1,35 @@ +/* + * Copyright (c) 2025, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ + +#include "cagra.cuh" +#include + +namespace cuvs::neighbors::cagra { + +#define RAFT_INST_CAGRA_MERGE(T, IdxT) \ + auto merge(raft::resources const& handle, \ + const cuvs::neighbors::cagra::merge_params& params, \ + std::vector*>& indices) \ + ->cuvs::neighbors::cagra::index \ + { \ + return cuvs::neighbors::cagra::merge(handle, params, indices); \ + } + +RAFT_INST_CAGRA_MERGE(uint8_t, uint32_t); + +#undef RAFT_INST_CAGRA_MERGE + +} // namespace cuvs::neighbors::cagra diff --git a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh new file mode 100644 index 0000000000..c429ecbed8 --- /dev/null +++ b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh @@ -0,0 +1,130 @@ +/* + * Copyright (c) 2023-2024, NVIDIA CORPORATION. + * + * Licensed under the Apache License, Version 2.0 (the "License"); + * you may not use this file except in compliance with the License. + * You may obtain a copy of the License at + * + * http://www.apache.org/licenses/LICENSE-2.0 + * + * Unless required by applicable law or agreed to in writing, software + * distributed under the License is distributed on an "AS IS" BASIS, + * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + * See the License for the specific language governing permissions and + * limitations under the License. + */ +#pragma once + +#include "../../../core/nvtx.hpp" +#include "../../vpq_dataset.cuh" +#include "cagra_build.cuh" +#include "graph_core.cuh" +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +#include + +// TODO: This shouldn't be calling spatial/knn APIs +#include "../ann_utils.cuh" + +#include + +#include +#include +#include + +namespace cuvs::neighbors::cagra::detail { + +template +index merge(raft::resources const& handle, + const cagra::merge_params& params, + std::vector*>& indices) +{ + std::size_t dim = 0; + std::size_t new_dataset_size = 0; + int64_t stride = -1; + + for (auto index : indices) { + using ds_idx_type = decltype(index->data().n_rows()); + if (auto* strided_dset = dynamic_cast*>(&index->data()); + strided_dset != nullptr) { + if (dim == 0) { + dim = index->dim(); + stride = strided_dset->stride(); + } else { + RAFT_EXPECTS(dim == index->dim(), "Dimension of datasets in indices must be equal."); + } + new_dataset_size += index->size(); + } else if (dynamic_cast*>(&index->data()) != + nullptr) { + RAFT_FAIL( + "cagra::merge only supports an index to which the dataset is attached. Please check if the " + "index was built with index_param.attach_dataset_on_build = true, or if a dataset was " + "attached after the build."); + } else { + RAFT_FAIL("cagra::merge only supports an uncompressed dataset index"); + } + } + + auto host_updated_dataset = raft::make_host_matrix(new_dataset_size, dim); + memset(host_updated_dataset.data_handle(), 0, sizeof(T) * host_updated_dataset.size()); + + IdxT offset = 0; + + for (auto index : indices) { + using ds_idx_type = decltype(index->data().n_rows()); + auto* strided_dset = dynamic_cast*>(&index->data()); + + RAFT_CUDA_TRY(cudaMemcpy2DAsync(host_updated_dataset.data_handle() + offset * dim, + sizeof(T) * dim, + strided_dset->view().data_handle(), + sizeof(T) * stride, + sizeof(T) * dim, + strided_dset->n_rows(), + cudaMemcpyDefault, + raft::resource::get_cuda_stream(handle))); + + offset += IdxT(index->data().n_rows()); + } + // Allocate the new dataset on device + auto device_updated_dataset = + raft::make_device_matrix(handle, new_dataset_size, dim); + auto device_updated_dataset_view = raft::make_device_matrix_view( + device_updated_dataset.data_handle(), new_dataset_size, dim); + + // Copy updated dataset on host memory to device memory + raft::copy(device_updated_dataset.data_handle(), + host_updated_dataset.data_handle(), + new_dataset_size * dim, + raft::resource::get_cuda_stream(handle)); + + auto merged_index = + cagra::build(handle, params, raft::make_const_mdspan(device_updated_dataset_view)); + + if (static_cast(stride) == dim) { + using out_mdarray_type = decltype(device_updated_dataset); + using out_layout_type = typename out_mdarray_type::layout_type; + using out_container_policy_type = typename out_mdarray_type::container_policy_type; + using out_owning_type = owning_dataset; + auto out_layout = raft::make_strided_layout(device_updated_dataset_view.extents(), + std::array{stride, 1}); + + merged_index.update_dataset(handle, + out_owning_type{std::move(device_updated_dataset), out_layout}); + } + return merged_index; +} + +} // namespace cuvs::neighbors::cagra::detail diff --git a/cpp/test/neighbors/ann_cagra.cuh b/cpp/test/neighbors/ann_cagra.cuh index bbafae6d18..1b8ff95c40 100644 --- a/cpp/test/neighbors/ann_cagra.cuh +++ b/cpp/test/neighbors/ann_cagra.cuh @@ -873,6 +873,190 @@ class AnnCagraFilterTest : public ::testing::TestWithParam { rmm::device_uvector search_queries; }; +template +class AnnCagraIndexMergeTest : public ::testing::TestWithParam { + public: + AnnCagraIndexMergeTest() + : stream_(raft::resource::get_cuda_stream(handle_)), + ps(::testing::TestWithParam::GetParam()), + database(0, stream_), + search_queries(0, stream_) + { + } + + protected: + void testCagra() + { + // TODO (rhdong): remove when NN Descent index building support InnerProduct. Reference + // issue: https://github.com/rapidsai/raft/issues/2276 + if (ps.metric == InnerProduct && ps.build_algo == graph_build_algo::NN_DESCENT) GTEST_SKIP(); + if (ps.compression != std::nullopt) GTEST_SKIP(); + + size_t queries_size = ps.n_queries * ps.k; + std::vector indices_Cagra(queries_size); + std::vector indices_naive(queries_size); + std::vector distances_Cagra(queries_size); + std::vector distances_naive(queries_size); + + { + rmm::device_uvector distances_naive_dev(queries_size, stream_); + rmm::device_uvector indices_naive_dev(queries_size, stream_); + + cuvs::neighbors::naive_knn(handle_, + distances_naive_dev.data(), + indices_naive_dev.data(), + search_queries.data(), + database.data(), + ps.n_queries, + ps.n_rows, + ps.dim, + ps.k, + ps.metric); + raft::update_host(distances_naive.data(), distances_naive_dev.data(), queries_size, stream_); + raft::update_host(indices_naive.data(), indices_naive_dev.data(), queries_size, stream_); + raft::resource::sync_stream(handle_); + } + + { + rmm::device_uvector distances_dev(queries_size, stream_); + rmm::device_uvector indices_dev(queries_size, stream_); + + { + cagra::index_params index_params; + index_params.metric = ps.metric; // Note: currently ony the cagra::index_params metric is + // not used for knn_graph building. + + switch (ps.build_algo) { + case graph_build_algo::IVF_PQ: + index_params.graph_build_params = + graph_build_params::ivf_pq_params(raft::matrix_extent(ps.n_rows, ps.dim)); + if (ps.ivf_pq_search_refine_ratio) { + std::get( + index_params.graph_build_params) + .refinement_rate = *ps.ivf_pq_search_refine_ratio; + } + break; + case graph_build_algo::NN_DESCENT: { + index_params.graph_build_params = + graph_build_params::nn_descent_params(index_params.intermediate_graph_degree); + break; + } + case graph_build_algo::AUTO: + // do nothing + break; + }; + + const double splite_ratio = 0.55; + const std::size_t database0_size = ps.n_rows * splite_ratio; + const std::size_t database1_size = ps.n_rows - database0_size; + + auto database0_view = raft::make_device_matrix_view( + (const DataT*)database.data(), database0_size, ps.dim); + + auto database1_view = raft::make_device_matrix_view( + (const DataT*)database.data() + database0_view.size(), database1_size, ps.dim); + + cagra::index index0(handle_); + cagra::index index1(handle_); + if (ps.host_dataset) { + { + std::optional> database_host{std::nullopt}; + database_host = raft::make_host_matrix(database0_size, ps.dim); + raft::copy(database_host->data_handle(), + database0_view.data_handle(), + database0_view.size(), + stream_); + auto database_host_view = raft::make_host_matrix_view( + (const DataT*)database_host->data_handle(), database0_size, ps.dim); + index0 = cagra::build(handle_, index_params, database_host_view); + } + { + std::optional> database_host{std::nullopt}; + database_host = raft::make_host_matrix(database1_size, ps.dim); + raft::copy(database_host->data_handle(), + database1_view.data_handle(), + database1_view.size(), + stream_); + auto database_host_view = raft::make_host_matrix_view( + (const DataT*)database_host->data_handle(), database1_size, ps.dim); + index1 = cagra::build(handle_, index_params, database_host_view); + } + } else { + index0 = cagra::build(handle_, index_params, database0_view); + index1 = cagra::build(handle_, index_params, database1_view); + }; + std::vector*> indices{&index0, &index1}; + cagra::merge_params merge_params{index_params}; + auto index = cagra::merge(handle_, merge_params, indices); + + auto search_queries_view = raft::make_device_matrix_view( + search_queries.data(), ps.n_queries, ps.dim); + auto indices_out_view = + raft::make_device_matrix_view(indices_dev.data(), ps.n_queries, ps.k); + auto dists_out_view = raft::make_device_matrix_view( + distances_dev.data(), ps.n_queries, ps.k); + + cagra::search_params search_params; + search_params.algo = ps.algo; + search_params.max_queries = ps.max_queries; + search_params.team_size = ps.team_size; + search_params.itopk_size = ps.itopk_size; + + cagra::search( + handle_, search_params, index, search_queries_view, indices_out_view, dists_out_view); + raft::update_host(distances_Cagra.data(), distances_dev.data(), queries_size, stream_); + raft::update_host(indices_Cagra.data(), indices_dev.data(), queries_size, stream_); + raft::resource::sync_stream(handle_); + } + + double min_recall = ps.min_recall; + EXPECT_TRUE(eval_neighbours(indices_naive, + indices_Cagra, + distances_naive, + distances_Cagra, + ps.n_queries, + ps.k, + 0.006, + min_recall)); + EXPECT_TRUE(eval_distances(handle_, + database.data(), + search_queries.data(), + indices_dev.data(), + distances_dev.data(), + ps.n_rows, + ps.dim, + ps.n_queries, + ps.k, + ps.metric, + 1.0e-4)); + } + } + + void SetUp() override + { + database.resize(((size_t)ps.n_rows) * ps.dim, stream_); + search_queries.resize(ps.n_queries * ps.dim, stream_); + raft::random::RngState r(1234ULL); + InitDataset(handle_, database.data(), ps.n_rows, ps.dim, ps.metric, r); + InitDataset(handle_, search_queries.data(), ps.n_queries, ps.dim, ps.metric, r); + raft::resource::sync_stream(handle_); + } + + void TearDown() override + { + raft::resource::sync_stream(handle_); + database.resize(0, stream_); + search_queries.resize(0, stream_); + } + + private: + raft::resources handle_; + rmm::cuda_stream_view stream_; + AnnCagraInputs ps; + rmm::device_uvector database; + rmm::device_uvector search_queries; +}; + inline std::vector generate_inputs() { // TODO(tfeher): test MULTI_CTA kernel with search_width > 1 to allow multiple CTA per queries diff --git a/cpp/test/neighbors/ann_cagra/test_float_uint32_t.cu b/cpp/test/neighbors/ann_cagra/test_float_uint32_t.cu index ca188d1320..e7bc6c7d4f 100644 --- a/cpp/test/neighbors/ann_cagra/test_float_uint32_t.cu +++ b/cpp/test/neighbors/ann_cagra/test_float_uint32_t.cu @@ -29,10 +29,16 @@ TEST_P(AnnCagraAddNodesTestF_U32, AnnCagraAddNodes) { this->testCagra(); } typedef AnnCagraFilterTest AnnCagraFilterTestF_U32; TEST_P(AnnCagraFilterTestF_U32, AnnCagra) { this->testCagra(); } +typedef AnnCagraIndexMergeTest AnnCagraIndexMergeTestF_U32; +TEST_P(AnnCagraIndexMergeTestF_U32, AnnCagraIndexMerge) { this->testCagra(); } + INSTANTIATE_TEST_CASE_P(AnnCagraTest, AnnCagraTestF_U32, ::testing::ValuesIn(inputs)); INSTANTIATE_TEST_CASE_P(AnnCagraAddNodesTest, AnnCagraAddNodesTestF_U32, ::testing::ValuesIn(inputs)); INSTANTIATE_TEST_CASE_P(AnnCagraFilterTest, AnnCagraFilterTestF_U32, ::testing::ValuesIn(inputs)); +INSTANTIATE_TEST_CASE_P(AnnCagraIndexMergeTest, + AnnCagraIndexMergeTestF_U32, + ::testing::ValuesIn(inputs)); } // namespace cuvs::neighbors::cagra diff --git a/cpp/test/neighbors/ann_cagra/test_half_uint32_t.cu b/cpp/test/neighbors/ann_cagra/test_half_uint32_t.cu index f03de69d2c..de682cee49 100644 --- a/cpp/test/neighbors/ann_cagra/test_half_uint32_t.cu +++ b/cpp/test/neighbors/ann_cagra/test_half_uint32_t.cu @@ -23,6 +23,12 @@ namespace cuvs::neighbors::cagra { typedef AnnCagraTest AnnCagraTestF16_U32; TEST_P(AnnCagraTestF16_U32, AnnCagra) { this->testCagra(); } +typedef AnnCagraIndexMergeTest AnnCagraIndexMergeTestF16_U32; +TEST_P(AnnCagraIndexMergeTestF16_U32, AnnCagraIndexMerge) { this->testCagra(); } + INSTANTIATE_TEST_CASE_P(AnnCagraTest, AnnCagraTestF16_U32, ::testing::ValuesIn(inputs)); +INSTANTIATE_TEST_CASE_P(AnnCagraIndexMergeTest, + AnnCagraIndexMergeTestF16_U32, + ::testing::ValuesIn(inputs)); } // namespace cuvs::neighbors::cagra diff --git a/cpp/test/neighbors/ann_cagra/test_int8_t_uint32_t.cu b/cpp/test/neighbors/ann_cagra/test_int8_t_uint32_t.cu index 4aa03afd58..0a8049490c 100644 --- a/cpp/test/neighbors/ann_cagra/test_int8_t_uint32_t.cu +++ b/cpp/test/neighbors/ann_cagra/test_int8_t_uint32_t.cu @@ -26,11 +26,15 @@ typedef AnnCagraAddNodesTest AnnCagraAddNodes TEST_P(AnnCagraAddNodesTestI8_U32, AnnCagra) { this->testCagra(); } typedef AnnCagraFilterTest AnnCagraFilterTestI8_U32; TEST_P(AnnCagraFilterTestI8_U32, AnnCagra) { this->testCagra(); } +typedef AnnCagraIndexMergeTest AnnCagraIndexMergeTestI8_U32; +TEST_P(AnnCagraIndexMergeTestI8_U32, AnnCagra) { this->testCagra(); } INSTANTIATE_TEST_CASE_P(AnnCagraTest, AnnCagraTestI8_U32, ::testing::ValuesIn(inputs)); INSTANTIATE_TEST_CASE_P(AnnCagraAddNodesTest, AnnCagraAddNodesTestI8_U32, ::testing::ValuesIn(inputs)); -INSTANTIATE_TEST_CASE_P(AnnCagraFilterTest, AnnCagraFilterTestI8_U32, ::testing::ValuesIn(inputs)); +INSTANTIATE_TEST_CASE_P(AnnCagraIndexMergeTest, + AnnCagraIndexMergeTestI8_U32, + ::testing::ValuesIn(inputs)); } // namespace cuvs::neighbors::cagra diff --git a/cpp/test/neighbors/ann_cagra/test_uint8_t_uint32_t.cu b/cpp/test/neighbors/ann_cagra/test_uint8_t_uint32_t.cu index b8e2a6b770..139723b2d3 100644 --- a/cpp/test/neighbors/ann_cagra/test_uint8_t_uint32_t.cu +++ b/cpp/test/neighbors/ann_cagra/test_uint8_t_uint32_t.cu @@ -26,11 +26,16 @@ typedef AnnCagraAddNodesTest AnnCagraAddNode TEST_P(AnnCagraAddNodesTestU8_U32, AnnCagra) { this->testCagra(); } typedef AnnCagraFilterTest AnnCagraFilterTestU8_U32; TEST_P(AnnCagraFilterTestU8_U32, AnnCagra) { this->testCagra(); } +typedef AnnCagraIndexMergeTest AnnCagraIndexMergeTestU8_U32; +TEST_P(AnnCagraIndexMergeTestU8_U32, AnnCagra) { this->testCagra(); } INSTANTIATE_TEST_CASE_P(AnnCagraTest, AnnCagraTestU8_U32, ::testing::ValuesIn(inputs)); INSTANTIATE_TEST_CASE_P(AnnCagraAddNodesTest, AnnCagraAddNodesTestU8_U32, ::testing::ValuesIn(inputs)); INSTANTIATE_TEST_CASE_P(AnnCagraFilterTest, AnnCagraFilterTestU8_U32, ::testing::ValuesIn(inputs)); +INSTANTIATE_TEST_CASE_P(AnnCagraIndexMergeTest, + AnnCagraIndexMergeTestU8_U32, + ::testing::ValuesIn(inputs)); } // namespace cuvs::neighbors::cagra From 696c6604639aeeefae819d05eac1dd5f518ea7be Mon Sep 17 00:00:00 2001 From: rhdong Date: Tue, 28 Jan 2025 14:09:22 -0800 Subject: [PATCH 02/10] simplify the memory logic --- .../neighbors/detail/cagra/cagra_merge.cuh | 20 ++++++------------- 1 file changed, 6 insertions(+), 14 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh index c429ecbed8..586585b509 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh @@ -78,8 +78,11 @@ index merge(raft::resources const& handle, } } - auto host_updated_dataset = raft::make_host_matrix(new_dataset_size, dim); - memset(host_updated_dataset.data_handle(), 0, sizeof(T) * host_updated_dataset.size()); + // Allocate the new dataset on device + auto device_updated_dataset = + raft::make_device_matrix(handle, new_dataset_size, dim); + auto device_updated_dataset_view = raft::make_device_matrix_view( + device_updated_dataset.data_handle(), new_dataset_size, dim); IdxT offset = 0; @@ -87,7 +90,7 @@ index merge(raft::resources const& handle, using ds_idx_type = decltype(index->data().n_rows()); auto* strided_dset = dynamic_cast*>(&index->data()); - RAFT_CUDA_TRY(cudaMemcpy2DAsync(host_updated_dataset.data_handle() + offset * dim, + RAFT_CUDA_TRY(cudaMemcpy2DAsync(device_updated_dataset.data_handle() + offset * dim, sizeof(T) * dim, strided_dset->view().data_handle(), sizeof(T) * stride, @@ -98,17 +101,6 @@ index merge(raft::resources const& handle, offset += IdxT(index->data().n_rows()); } - // Allocate the new dataset on device - auto device_updated_dataset = - raft::make_device_matrix(handle, new_dataset_size, dim); - auto device_updated_dataset_view = raft::make_device_matrix_view( - device_updated_dataset.data_handle(), new_dataset_size, dim); - - // Copy updated dataset on host memory to device memory - raft::copy(device_updated_dataset.data_handle(), - host_updated_dataset.data_handle(), - new_dataset_size * dim, - raft::resource::get_cuda_stream(handle)); auto merged_index = cagra::build(handle, params, raft::make_const_mdspan(device_updated_dataset_view)); From e5067c889981257ed5f2aa264b4565eba3024f70 Mon Sep 17 00:00:00 2001 From: rhdong Date: Wed, 29 Jan 2025 19:50:28 -0800 Subject: [PATCH 03/10] automatically the memory choose & judgment owning for better readability --- cpp/include/cuvs/neighbors/cagra.hpp | 22 ++++- .../neighbors/detail/cagra/cagra_merge.cuh | 99 ++++++++++++------- 2 files changed, 80 insertions(+), 41 deletions(-) diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index 8f257900ea..8a4757c821 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -304,7 +304,7 @@ struct index : cuvs::neighbors::index { return data_rows > 0 ? data_rows : graph_view_.extent(0); } - /** Dimensionality of the data. */ + /** dimension of the data. */ [[nodiscard]] constexpr inline auto dim() const noexcept -> uint32_t { return dataset_->dim(); } /** Graph degree */ [[nodiscard]] constexpr inline auto graph_degree() const noexcept -> uint32_t @@ -1772,6 +1772,9 @@ void serialize_to_hnswlib(raft::resources const& handle, * This function merges multiple CAGRA indices into one, combining both the datasets and graph * structures. * + * @note: When device memory is sufficient, the dataset attached to the returned index is allocated + * in device memory by default; otherwise, host memory is used automatically. + * * Usage example: * @code{.cpp} * using namespace raft::neighbors; @@ -1790,7 +1793,7 @@ void serialize_to_hnswlib(raft::resources const& handle, * @param[in] res RAFT resources used for the merge operation. * @param[in] params Parameters that control the merging process. * @param[in] indices A vector of pointers to the CAGRA indices to merge. All indices must: - * - Have attached datasets with the same dimensionality. + * - Have attached datasets with the same dimension. * * @return A new CAGRA index containing the merged indices, graph, and dataset. */ @@ -1804,6 +1807,9 @@ auto merge(raft::resources const& res, * This function merges multiple CAGRA indices into one, combining both the datasets and graph * structures. * + * @note: When device memory is sufficient, the dataset attached to the returned index is allocated + * in device memory by default; otherwise, host memory is used automatically. + * * Usage example: * @code{.cpp} * using namespace raft::neighbors; @@ -1822,7 +1828,7 @@ auto merge(raft::resources const& res, * @param[in] res RAFT resources used for the merge operation. * @param[in] params Parameters that control the merging process. * @param[in] indices A vector of pointers to the CAGRA indices to merge. All indices must: - * - Have attached datasets with the same dimensionality. + * - Have attached datasets with the same dimension. * * @return A new CAGRA index containing the merged indices, graph, and dataset. */ @@ -1836,6 +1842,9 @@ auto merge(raft::resources const& res, * This function merges multiple CAGRA indices into one, combining both the datasets and graph * structures. * + * @note: When device memory is sufficient, the dataset attached to the returned index is allocated + * in device memory by default; otherwise, host memory is used automatically. + * * Usage example: * @code{.cpp} * using namespace raft::neighbors; @@ -1854,7 +1863,7 @@ auto merge(raft::resources const& res, * @param[in] res RAFT resources used for the merge operation. * @param[in] params Parameters that control the merging process. * @param[in] indices A vector of pointers to the CAGRA indices to merge. All indices must: - * - Have attached datasets with the same dimensionality. + * - Have attached datasets with the same dimension. * * @return A new CAGRA index containing the merged indices, graph, and dataset. */ @@ -1868,6 +1877,9 @@ auto merge(raft::resources const& res, * This function merges multiple CAGRA indices into one, combining both the datasets and graph * structures. * + * @note: When device memory is sufficient, the dataset attached to the returned index is allocated + * in device memory by default; otherwise, host memory is used automatically. + * * Usage example: * @code{.cpp} * using namespace raft::neighbors; @@ -1886,7 +1898,7 @@ auto merge(raft::resources const& res, * @param[in] res RAFT resources used for the merge operation. * @param[in] params Parameters that control the merging process. * @param[in] indices A vector of pointers to the CAGRA indices to merge. All indices must: - * - Have attached datasets with the same dimensionality. + * - Have attached datasets with the same dimension. * * @return A new CAGRA index containing the merged indices, graph, and dataset. */ diff --git a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh index 586585b509..8f85450d6a 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh @@ -78,45 +78,72 @@ index merge(raft::resources const& handle, } } - // Allocate the new dataset on device - auto device_updated_dataset = - raft::make_device_matrix(handle, new_dataset_size, dim); - auto device_updated_dataset_view = raft::make_device_matrix_view( - device_updated_dataset.data_handle(), new_dataset_size, dim); - IdxT offset = 0; - for (auto index : indices) { - using ds_idx_type = decltype(index->data().n_rows()); - auto* strided_dset = dynamic_cast*>(&index->data()); - - RAFT_CUDA_TRY(cudaMemcpy2DAsync(device_updated_dataset.data_handle() + offset * dim, - sizeof(T) * dim, - strided_dset->view().data_handle(), - sizeof(T) * stride, - sizeof(T) * dim, - strided_dset->n_rows(), - cudaMemcpyDefault, - raft::resource::get_cuda_stream(handle))); - - offset += IdxT(index->data().n_rows()); - } - - auto merged_index = - cagra::build(handle, params, raft::make_const_mdspan(device_updated_dataset_view)); - - if (static_cast(stride) == dim) { - using out_mdarray_type = decltype(device_updated_dataset); - using out_layout_type = typename out_mdarray_type::layout_type; - using out_container_policy_type = typename out_mdarray_type::container_policy_type; - using out_owning_type = owning_dataset; - auto out_layout = raft::make_strided_layout(device_updated_dataset_view.extents(), - std::array{stride, 1}); - - merged_index.update_dataset(handle, - out_owning_type{std::move(device_updated_dataset), out_layout}); + // Allocate the new dataset on device + bool dataset_on_device = cuvs::neighbors::nn_descent::has_enough_device_memory( + handle, raft::make_extents(new_dataset_size, dim), sizeof(IdxT)); + + auto merge_dataset = [&](T* dst) { + for (auto index : indices) { + using ds_idx_type = decltype(index->data().n_rows()); + auto* strided_dset = dynamic_cast*>(&index->data()); + + RAFT_CUDA_TRY(cudaMemcpy2DAsync(dst + offset * dim, + sizeof(T) * dim, + strided_dset->view().data_handle(), + sizeof(T) * stride, + sizeof(T) * dim, + strided_dset->n_rows(), + cudaMemcpyDefault, + raft::resource::get_cuda_stream(handle))); + + offset += IdxT(index->data().n_rows()); + } + }; + + if (dataset_on_device) { + RAFT_LOG_DEBUG("cagra merge: using device memory for merged dataset"); + + auto updated_dataset = raft::make_device_matrix( + handle, std::int64_t(new_dataset_size), std::int64_t(dim)); + + merge_dataset(updated_dataset.data_handle()); + + auto merged_index = + cagra::build(handle, params, raft::make_const_mdspan(updated_dataset.view())); + if (!merged_index.data().is_owning() && params.attach_dataset_on_build) { + using matrix_t = decltype(updated_dataset); + using layout_t = typename matrix_t::layout_type; + using container_policy_t = typename matrix_t::container_policy_type; + using owning_t = owning_dataset; + auto out_layout = raft::make_strided_layout(updated_dataset.view().extents(), + std::array{stride, 1}); + merged_index.update_dataset(handle, owning_t{std::move(updated_dataset), out_layout}); + } + return merged_index; + + } else { + RAFT_LOG_DEBUG("cagra::merge: using host memory for merged dataset"); + + auto updated_dataset = + raft::make_host_matrix(std::int64_t(new_dataset_size), std::int64_t(dim)); + + merge_dataset(updated_dataset.data_handle()); + + auto merged_index = + cagra::build(handle, params, raft::make_const_mdspan(updated_dataset.view())); + if (!merged_index.data().is_owning() && params.attach_dataset_on_build) { + using matrix_t = decltype(updated_dataset); + using layout_t = typename matrix_t::layout_type; + using container_policy_t = typename matrix_t::container_policy_type; + using owning_t = owning_dataset; + auto out_layout = raft::make_strided_layout(updated_dataset.view().extents(), + std::array{stride, 1}); + merged_index.update_dataset(handle, owning_t{std::move(updated_dataset), out_layout}); + } + return merged_index; } - return merged_index; } } // namespace cuvs::neighbors::cagra::detail From 8b95d7bd5a61e2c785bd19bac431c299b12d2ea8 Mon Sep 17 00:00:00 2001 From: rhdong Date: Fri, 31 Jan 2025 11:49:23 -0800 Subject: [PATCH 04/10] add `nullptr` checking & reduce binary size --- cpp/src/neighbors/detail/cagra/cagra_merge.cuh | 9 ++------- 1 file changed, 2 insertions(+), 7 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh index 8f85450d6a..a854a414a5 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh @@ -15,10 +15,6 @@ */ #pragma once -#include "../../../core/nvtx.hpp" -#include "../../vpq_dataset.cuh" -#include "cagra_build.cuh" -#include "graph_core.cuh" #include #include @@ -36,9 +32,6 @@ #include -// TODO: This shouldn't be calling spatial/knn APIs -#include "../ann_utils.cuh" - #include #include @@ -57,6 +50,8 @@ index merge(raft::resources const& handle, int64_t stride = -1; for (auto index : indices) { + RAFT_EXPECTS(index != nullptr, + "Null pointer detected in 'indices'. Ensure all elements are valid before usage."); using ds_idx_type = decltype(index->data().n_rows()); if (auto* strided_dset = dynamic_cast*>(&index->data()); strided_dset != nullptr) { From 7446f0ed635e31e9c568e1f86aceb08907d8a6a4 Mon Sep 17 00:00:00 2001 From: rhdong Date: Fri, 31 Jan 2025 11:54:12 -0800 Subject: [PATCH 05/10] increase max_allowed_size_compressed to '1.2G' with https://github.com/rapidsai/cuvs/issues/634 --- python/libcuvs/pyproject.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/python/libcuvs/pyproject.toml b/python/libcuvs/pyproject.toml index 28443b7820..3fd69b9899 100644 --- a/python/libcuvs/pyproject.toml +++ b/python/libcuvs/pyproject.toml @@ -105,4 +105,4 @@ select = [ ] # detect when package size grows significantly -max_allowed_size_compressed = '1.1G' +max_allowed_size_compressed = '1.2G' From 690e7750d889bb920ad2cb96404d1416aa27033e Mon Sep 17 00:00:00 2001 From: rhdong Date: Sun, 2 Feb 2025 15:04:01 -0800 Subject: [PATCH 06/10] decouple the `merge_params` from `index_params` --- cpp/include/cuvs/neighbors/cagra.hpp | 9 +++++++-- cpp/src/neighbors/detail/cagra/cagra_merge.cuh | 10 ++++++---- cpp/tests/neighbors/ann_cagra.cuh | 4 ++++ 3 files changed, 17 insertions(+), 6 deletions(-) diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index 72d4fde5aa..3b840496a4 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -272,8 +272,13 @@ struct extend_params { * @defgroup cagra_cpp_merge_params CAGRA index merge parameters * @{ */ -struct merge_params : public index_params { - merge_params(const index_params& params) : index_params(params) {} +struct merge_params { + merge_params() = default; + + explicit merge_params(const cagra::index_params& params) : output_index_params(params) {} + + // Parameters for creating the output index + cagra::index_params output_index_params; }; /** diff --git a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh index a854a414a5..9bc1456411 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh @@ -97,6 +97,8 @@ index merge(raft::resources const& handle, } }; + cagra::index_params output_index_params = params.output_index_params; + if (dataset_on_device) { RAFT_LOG_DEBUG("cagra merge: using device memory for merged dataset"); @@ -106,8 +108,8 @@ index merge(raft::resources const& handle, merge_dataset(updated_dataset.data_handle()); auto merged_index = - cagra::build(handle, params, raft::make_const_mdspan(updated_dataset.view())); - if (!merged_index.data().is_owning() && params.attach_dataset_on_build) { + cagra::build(handle, output_index_params, raft::make_const_mdspan(updated_dataset.view())); + if (!merged_index.data().is_owning() && output_index_params.attach_dataset_on_build) { using matrix_t = decltype(updated_dataset); using layout_t = typename matrix_t::layout_type; using container_policy_t = typename matrix_t::container_policy_type; @@ -127,8 +129,8 @@ index merge(raft::resources const& handle, merge_dataset(updated_dataset.data_handle()); auto merged_index = - cagra::build(handle, params, raft::make_const_mdspan(updated_dataset.view())); - if (!merged_index.data().is_owning() && params.attach_dataset_on_build) { + cagra::build(handle, output_index_params, raft::make_const_mdspan(updated_dataset.view())); + if (!merged_index.data().is_owning() && output_index_params.attach_dataset_on_build) { using matrix_t = decltype(updated_dataset); using layout_t = typename matrix_t::layout_type; using container_policy_t = typename matrix_t::container_policy_type; diff --git a/cpp/tests/neighbors/ann_cagra.cuh b/cpp/tests/neighbors/ann_cagra.cuh index 7aa31a1666..e1eba5a9df 100644 --- a/cpp/tests/neighbors/ann_cagra.cuh +++ b/cpp/tests/neighbors/ann_cagra.cuh @@ -955,6 +955,10 @@ class AnnCagraIndexMergeTest : public ::testing::TestWithParam { graph_build_params::nn_descent_params(index_params.intermediate_graph_degree); break; } + case graph_build_algo::ITERATIVE_CAGRA_SEARCH: { + index_params.graph_build_params = graph_build_params::iterative_search_params(); + break; + } case graph_build_algo::AUTO: // do nothing break; From afb6026c4cae0555c02ad0c380ce45f9b90951c6 Mon Sep 17 00:00:00 2001 From: rhdong Date: Tue, 4 Feb 2025 08:03:25 -0800 Subject: [PATCH 07/10] fix device memory allocation strategy --- cpp/src/neighbors/detail/cagra/cagra_merge.cuh | 13 +++---------- cpp/tests/neighbors/ann_cagra.cuh | 13 ++++++++++++- 2 files changed, 15 insertions(+), 11 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh index 9bc1456411..bc29cb2060 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh @@ -30,8 +30,6 @@ #include #include -#include - #include #include @@ -75,10 +73,6 @@ index merge(raft::resources const& handle, IdxT offset = 0; - // Allocate the new dataset on device - bool dataset_on_device = cuvs::neighbors::nn_descent::has_enough_device_memory( - handle, raft::make_extents(new_dataset_size, dim), sizeof(IdxT)); - auto merge_dataset = [&](T* dst) { for (auto index : indices) { using ds_idx_type = decltype(index->data().n_rows()); @@ -99,9 +93,7 @@ index merge(raft::resources const& handle, cagra::index_params output_index_params = params.output_index_params; - if (dataset_on_device) { - RAFT_LOG_DEBUG("cagra merge: using device memory for merged dataset"); - + try { auto updated_dataset = raft::make_device_matrix( handle, std::int64_t(new_dataset_size), std::int64_t(dim)); @@ -118,9 +110,10 @@ index merge(raft::resources const& handle, std::array{stride, 1}); merged_index.update_dataset(handle, owning_t{std::move(updated_dataset), out_layout}); } + RAFT_LOG_DEBUG("cagra merge: using device memory for merged dataset"); return merged_index; - } else { + } catch (std::bad_alloc& e) { RAFT_LOG_DEBUG("cagra::merge: using host memory for merged dataset"); auto updated_dataset = diff --git a/cpp/tests/neighbors/ann_cagra.cuh b/cpp/tests/neighbors/ann_cagra.cuh index 1e009f7162..1e695f9a8f 100644 --- a/cpp/tests/neighbors/ann_cagra.cuh +++ b/cpp/tests/neighbors/ann_cagra.cuh @@ -936,10 +936,21 @@ class AnnCagraIndexMergeTest : public ::testing::TestWithParam { protected: void testCagra() { - // TODO (rhdong): remove when NN Descent index building support InnerProduct. Reference + // TODO (tarang-jain): remove when NN Descent index building support InnerProduct. Reference // issue: https://github.com/rapidsai/raft/issues/2276 if (ps.metric == InnerProduct && ps.build_algo == graph_build_algo::NN_DESCENT) GTEST_SKIP(); if (ps.compression != std::nullopt) GTEST_SKIP(); + // IVF_PQ and NN_DESCENT graph builds do not support BitwiseHamming + if (ps.metric == cuvs::distance::DistanceType::BitwiseHamming && + ((!std::is_same_v) || + (ps.build_algo != graph_build_algo::ITERATIVE_CAGRA_SEARCH))) + GTEST_SKIP(); + // If the dataset dimension is small and the dataset size is large, there can be a lot of + // dataset vectors that have the same distance to the query, especially in the binary Hamming + // distance, making it impossible to make a top-k ground truth. + if (ps.metric == cuvs::distance::DistanceType::BitwiseHamming && + (ps.k * ps.dim * 8 / 5 /*(=magic number)*/ < ps.n_rows)) + GTEST_SKIP(); size_t queries_size = ps.n_queries * ps.k; std::vector indices_Cagra(queries_size); From da45bdb69baad0e29f11b77be50cd9c0460eeb94 Mon Sep 17 00:00:00 2001 From: rhdong Date: Wed, 5 Feb 2025 13:50:28 -0800 Subject: [PATCH 08/10] add merge strategy for future extension --- cpp/include/cuvs/neighbors/cagra.hpp | 54 ++++++++++++++++++++++++++-- 1 file changed, 52 insertions(+), 2 deletions(-) diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index 3b840496a4..4541df7260 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -272,13 +272,63 @@ struct extend_params { * @defgroup cagra_cpp_merge_params CAGRA index merge parameters * @{ */ + +/** + * @brief Determines the strategy for merging CAGRA graphs. + * + * @note Currently, only the PHYSICAL strategy is supported. + */ +enum MergeStrategy { + /** + * @brief Physical merge: Builds a new CAGRA graph from the union of dataset points + * in existing CAGRA graphs. + * + * This is expensive to build but does not impact search latency or quality. + * Preferred for many smaller CAGRA graphs. + * + * @note Currently, this is the only supported strategy. + */ + PHYSICAL, + + /** + * @brief Logical merge: Wraps a new index structure around existing CAGRA graphs + * and broadcasts the query to each of them. + * + * This is a fast merge but incurs a small hit in search latency. + * Preferred for fewer larger CAGRA graphs. + * + * @note Not supported yet. + */ + LOGICAL, + + /** + * @brief Smart merge: Overlaps dataset vectors across CAGRA graphs and merges + * the graphs into a single graph. + * + * This is suitable for many larger CAGRA graphs. + * + * @note Not supported yet. + */ + SMART +}; + +/** + * @brief Parameters for merging CAGRA indexes. + */ struct merge_params { merge_params() = default; + /** + * @brief Constructs merge parameters with given index parameters. + * @param params Parameters for creating the output index. + */ explicit merge_params(const cagra::index_params& params) : output_index_params(params) {} - // Parameters for creating the output index + /// Parameters for creating the output index. cagra::index_params output_index_params; + + /// Strategy for merging. Defaults to `MergeStrategy::PHYSICAL`. + MergeStrategy strategy = MergeStrategy::PHYSICAL; }; /** @@ -325,7 +375,7 @@ struct index : cuvs::neighbors::index { return data_rows > 0 ? data_rows : graph_view_.extent(0); } - /** dimension of the data. */ + /** Dimensionality of the data. */ [[nodiscard]] constexpr inline auto dim() const noexcept -> uint32_t { return dataset_->dim(); } /** Graph degree */ [[nodiscard]] constexpr inline auto graph_degree() const noexcept -> uint32_t From 9d3acb04fe5ebea8e4f1bab8de9bda90d0f3235a Mon Sep 17 00:00:00 2001 From: rhdong Date: Thu, 6 Feb 2025 09:17:27 -0800 Subject: [PATCH 09/10] reserve `PHYSICAL` only --- cpp/include/cuvs/neighbors/cagra.hpp | 23 +------------------ .../neighbors/detail/cagra/cagra_merge.cuh | 4 ++-- 2 files changed, 3 insertions(+), 24 deletions(-) diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index 4541df7260..1d0acbe355 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -288,28 +288,7 @@ enum MergeStrategy { * * @note Currently, this is the only supported strategy. */ - PHYSICAL, - - /** - * @brief Logical merge: Wraps a new index structure around existing CAGRA graphs - * and broadcasts the query to each of them. - * - * This is a fast merge but incurs a small hit in search latency. - * Preferred for fewer larger CAGRA graphs. - * - * @note Not supported yet. - */ - LOGICAL, - - /** - * @brief Smart merge: Overlaps dataset vectors across CAGRA graphs and merges - * the graphs into a single graph. - * - * This is suitable for many larger CAGRA graphs. - * - * @note Not supported yet. - */ - SMART + PHYSICAL }; /** diff --git a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh index bc29cb2060..39b99f9312 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh @@ -111,7 +111,7 @@ index merge(raft::resources const& handle, merged_index.update_dataset(handle, owning_t{std::move(updated_dataset), out_layout}); } RAFT_LOG_DEBUG("cagra merge: using device memory for merged dataset"); - return merged_index; + return std::move(merged_index); } catch (std::bad_alloc& e) { RAFT_LOG_DEBUG("cagra::merge: using host memory for merged dataset"); @@ -132,7 +132,7 @@ index merge(raft::resources const& handle, std::array{stride, 1}); merged_index.update_dataset(handle, owning_t{std::move(updated_dataset), out_layout}); } - return merged_index; + return std::move(merged_index); } } From ea7991b8afa64863b9d0a06245c4473b08a1dd3f Mon Sep 17 00:00:00 2001 From: rhdong Date: Thu, 6 Feb 2025 09:31:02 -0800 Subject: [PATCH 10/10] revert the std::move --- cpp/src/neighbors/detail/cagra/cagra_merge.cuh | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh index 39b99f9312..bc29cb2060 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh @@ -111,7 +111,7 @@ index merge(raft::resources const& handle, merged_index.update_dataset(handle, owning_t{std::move(updated_dataset), out_layout}); } RAFT_LOG_DEBUG("cagra merge: using device memory for merged dataset"); - return std::move(merged_index); + return merged_index; } catch (std::bad_alloc& e) { RAFT_LOG_DEBUG("cagra::merge: using host memory for merged dataset"); @@ -132,7 +132,7 @@ index merge(raft::resources const& handle, std::array{stride, 1}); merged_index.update_dataset(handle, owning_t{std::move(updated_dataset), out_layout}); } - return std::move(merged_index); + return merged_index; } }