diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index b324c6786..6cd87cbad 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -396,6 +396,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 d8edf742b..1d0acbe35 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -264,6 +264,51 @@ struct extend_params { * 0. */ uint32_t max_chunk_size = 0; }; +/** + * @} + */ + +/** + * @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 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. + cagra::index_params output_index_params; + + /// Strategy for merging. Defaults to `MergeStrategy::PHYSICAL`. + MergeStrategy strategy = MergeStrategy::PHYSICAL; +}; /** * @} @@ -1794,6 +1839,150 @@ void serialize_to_hnswlib( std::optional> dataset = std::nullopt); +/** + * @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. + * + * @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; + * 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 dimension. + * + * @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. + * + * @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; + * 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 dimension. + * + * @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. + * + * @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; + * 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 dimension. + * + * @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. + * + * @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; + * 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 dimension. + * + * @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 f294c9b44..305788f7c 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" @@ -380,6 +381,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 000000000..951c0c5fe --- /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 000000000..704a00f74 --- /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 000000000..a7e903562 --- /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 000000000..a4fc7149c --- /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 000000000..bc29cb206 --- /dev/null +++ b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh @@ -0,0 +1,139 @@ +/* + * 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 + +#include +#include +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +#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) { + 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) { + 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"); + } + } + + IdxT offset = 0; + + 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()); + } + }; + + cagra::index_params output_index_params = params.output_index_params; + + try { + 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, 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; + 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}); + } + RAFT_LOG_DEBUG("cagra merge: using device memory for merged dataset"); + return merged_index; + + } catch (std::bad_alloc& e) { + 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, 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; + 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; + } +} + +} // namespace cuvs::neighbors::cagra::detail diff --git a/cpp/tests/neighbors/ann_cagra.cuh b/cpp/tests/neighbors/ann_cagra.cuh index aedb11543..1e695f9a8 100644 --- a/cpp/tests/neighbors/ann_cagra.cuh +++ b/cpp/tests/neighbors/ann_cagra.cuh @@ -922,6 +922,205 @@ 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 (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); + 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::ITERATIVE_CAGRA_SEARCH: { + index_params.graph_build_params = graph_build_params::iterative_search_params(); + 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/tests/neighbors/ann_cagra/test_float_uint32_t.cu b/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu index 640657ccb..16c0d3d2a 100644 --- a/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu +++ b/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu @@ -29,6 +29,9 @@ 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, @@ -36,5 +39,8 @@ INSTANTIATE_TEST_CASE_P(AnnCagraAddNodesTest, INSTANTIATE_TEST_CASE_P(AnnCagraFilterTest, AnnCagraFilterTestF_U32, ::testing::ValuesIn(inputs_filtering)); +INSTANTIATE_TEST_CASE_P(AnnCagraIndexMergeTest, + AnnCagraIndexMergeTestF_U32, + ::testing::ValuesIn(inputs)); } // namespace cuvs::neighbors::cagra diff --git a/cpp/tests/neighbors/ann_cagra/test_half_uint32_t.cu b/cpp/tests/neighbors/ann_cagra/test_half_uint32_t.cu index f03de69d2..de682cee4 100644 --- a/cpp/tests/neighbors/ann_cagra/test_half_uint32_t.cu +++ b/cpp/tests/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/tests/neighbors/ann_cagra/test_int8_t_uint32_t.cu b/cpp/tests/neighbors/ann_cagra/test_int8_t_uint32_t.cu index 32f6289fe..28be13c5b 100644 --- a/cpp/tests/neighbors/ann_cagra/test_int8_t_uint32_t.cu +++ b/cpp/tests/neighbors/ann_cagra/test_int8_t_uint32_t.cu @@ -26,6 +26,8 @@ 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, @@ -34,5 +36,8 @@ INSTANTIATE_TEST_CASE_P(AnnCagraAddNodesTest, INSTANTIATE_TEST_CASE_P(AnnCagraFilterTest, AnnCagraFilterTestI8_U32, ::testing::ValuesIn(inputs_filtering)); +INSTANTIATE_TEST_CASE_P(AnnCagraIndexMergeTest, + AnnCagraIndexMergeTestI8_U32, + ::testing::ValuesIn(inputs)); } // namespace cuvs::neighbors::cagra diff --git a/cpp/tests/neighbors/ann_cagra/test_uint8_t_uint32_t.cu b/cpp/tests/neighbors/ann_cagra/test_uint8_t_uint32_t.cu index 53f804be6..dd9ffd1ae 100644 --- a/cpp/tests/neighbors/ann_cagra/test_uint8_t_uint32_t.cu +++ b/cpp/tests/neighbors/ann_cagra/test_uint8_t_uint32_t.cu @@ -26,6 +26,8 @@ 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, @@ -34,5 +36,8 @@ INSTANTIATE_TEST_CASE_P(AnnCagraAddNodesTest, INSTANTIATE_TEST_CASE_P(AnnCagraFilterTest, AnnCagraFilterTestU8_U32, ::testing::ValuesIn(inputs_filtering)); +INSTANTIATE_TEST_CASE_P(AnnCagraIndexMergeTest, + AnnCagraIndexMergeTestU8_U32, + ::testing::ValuesIn(inputs)); } // namespace cuvs::neighbors::cagra diff --git a/python/libcuvs/pyproject.toml b/python/libcuvs/pyproject.toml index 28443b782..3fd69b989 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'