diff --git a/c/include/cuvs/neighbors/cagra.h b/c/include/cuvs/neighbors/cagra.h index ed8c69c830..16e92b1124 100644 --- a/c/include/cuvs/neighbors/cagra.h +++ b/c/include/cuvs/neighbors/cagra.h @@ -418,33 +418,6 @@ cuvsError_t cuvsCagraIndexGetDataset(cuvsCagraIndex_t index, DLManagedTensor* da */ cuvsError_t cuvsCagraIndexGetGraph(cuvsCagraIndex_t index, DLManagedTensor* graph); -/** - * @} - */ - -/** - * @defgroup cagra_c_merge_params C API for CUDA ANN Graph-based nearest neighbor search - * @{ - */ - -/** - * @brief Supplemental parameters to merge CAGRA index - * - */ - -struct cuvsCagraMergeParams { - cuvsCagraIndexParams_t output_index_params; - cuvsMergeStrategy strategy; -}; - -typedef struct cuvsCagraMergeParams* cuvsCagraMergeParams_t; - -/** Allocate CAGRA merge params with default values */ -cuvsError_t cuvsCagraMergeParamsCreate(cuvsCagraMergeParams_t* params); - -/** De-allocate CAGRA merge params */ -cuvsError_t cuvsCagraMergeParamsDestroy(cuvsCagraMergeParams_t params); - /** * @} */ @@ -718,7 +691,7 @@ cuvsError_t cuvsCagraIndexFromArgs(cuvsResources_t res, * * All input indices must have been built with the same data type (`index.dtype`) and * have the same dimensionality (`index.dims`). The merged index uses the output - * parameters specified in `cuvsCagraMergeParams`. + * parameters specified in `cuvsCagraIndexParams`. * * Input indices must have: * - `index.dtype.code` and `index.dtype.bits` matching across all indices. @@ -745,29 +718,31 @@ cuvsError_t cuvsCagraIndexFromArgs(cuvsResources_t res, * * // Assume index1 and index2 have been built using cuvsCagraBuild * - * cuvsCagraMergeParams_t merge_params; - * cuvsError_t params_create_status = cuvsCagraMergeParamsCreate(&merge_params); + * cuvsCagraIndexParams_t merge_params; + * cuvsError_t params_create_status = cuvsCagraIndexParamsCreate(&merge_params); * * cuvsError_t merge_status = cuvsCagraMerge(res, merge_params, (cuvsCagraIndex_t[]){index1, * index2}, 2, merged_index); * * // Use merged_index for search operations * - * cuvsError_t params_destroy_status = cuvsCagraMergeParamsDestroy(merge_params); + * cuvsError_t params_destroy_status = cuvsCagraIndexParamsDestroy(merge_params); * cuvsError_t res_destroy_status = cuvsResourcesDestroy(res); * @endcode * * @param[in] res cuvsResources_t opaque C handle - * @param[in] params cuvsCagraMergeParams_t parameters controlling merge behavior + * @param[in] params cuvsCagraIndexParams_t parameters controlling merge behavior * @param[in] indices Array of input cuvsCagraIndex_t handles to merge * @param[in] num_indices Number of input indices + * @param[in] filter Filter that can be used to filter out vectors from the merged index * @param[out] output_index Output handle that will store the merged index. * Must be initialized using `cuvsCagraIndexCreate` before use. */ cuvsError_t cuvsCagraMerge(cuvsResources_t res, - cuvsCagraMergeParams_t params, + cuvsCagraIndexParams_t params, cuvsCagraIndex_t* indices, size_t num_indices, + cuvsFilter filter, cuvsCagraIndex_t output_index); /** diff --git a/c/src/neighbors/cagra.cpp b/c/src/neighbors/cagra.cpp index 8df6e387b3..cb8699db79 100644 --- a/c/src/neighbors/cagra.cpp +++ b/c/src/neighbors/cagra.cpp @@ -285,23 +285,23 @@ void* _deserialize(cuvsResources_t res, const char* filename) template void* _merge(cuvsResources_t res, - cuvsCagraMergeParams params, + cuvsCagraIndexParams params, cuvsCagraIndex_t* indices, - size_t num_indices) + size_t num_indices, + cuvsFilter filter) { auto res_ptr = reinterpret_cast(res); - cuvs::neighbors::cagra::merge_params merge_params_cpp; - auto& out_idx_params = *params.output_index_params; + cuvs::neighbors::cagra::index_params params_cpp; - merge_params_cpp.output_index_params.metric = - static_cast((int)out_idx_params.metric); - merge_params_cpp.output_index_params.intermediate_graph_degree = - out_idx_params.intermediate_graph_degree; - merge_params_cpp.output_index_params.graph_degree = out_idx_params.graph_degree; + params_cpp.metric = + static_cast((int)params.metric); + params_cpp.intermediate_graph_degree = + params.intermediate_graph_degree; + params_cpp.graph_degree = params.graph_degree; int64_t total_size = 0; int64_t dim = 0; - if (out_idx_params.build_algo == cuvsCagraGraphBuildAlgo::IVF_PQ) { + if (params.build_algo == cuvsCagraGraphBuildAlgo::IVF_PQ) { auto first_idx_ptr = reinterpret_cast*>(indices[0]->addr); dim = first_idx_ptr->dim(); @@ -312,9 +312,9 @@ void* _merge(cuvsResources_t res, } } - _set_graph_build_params(merge_params_cpp.output_index_params.graph_build_params, - out_idx_params, - out_idx_params.build_algo, + _set_graph_build_params(params_cpp.graph_build_params, + params, + params.build_algo, total_size, dim); @@ -325,10 +325,21 @@ void* _merge(cuvsResources_t res, index_ptrs.push_back(idx_ptr); } - auto merged_index = new cuvs::neighbors::cagra::index( - cuvs::neighbors::cagra::merge(*res_ptr, merge_params_cpp, index_ptrs)); - - return merged_index; + if (filter.type == NO_FILTER) { + return new cuvs::neighbors::cagra::index( + cuvs::neighbors::cagra::merge(*res_ptr, params_cpp, index_ptrs)); + } else if (filter.type == BITSET) { + using filter_mdspan_type = raft::device_vector_view; + auto removed_indices_tensor = reinterpret_cast(filter.addr); + auto removed_indices = cuvs::core::from_dlpack(removed_indices_tensor); + cuvs::core::bitset_view removed_indices_bitset( + removed_indices, total_size); + auto bitset_filter_obj = cuvs::neighbors::filtering::bitset_filter(removed_indices_bitset); + return new cuvs::neighbors::cagra::index( + cuvs::neighbors::cagra::merge(*res_ptr, params_cpp, index_ptrs, bitset_filter_obj)); + } else { + RAFT_FAIL("Unsupported filter type: BITMAP"); + } } template @@ -601,9 +612,10 @@ extern "C" cuvsError_t cuvsCagraSearch(cuvsResources_t res, } extern "C" cuvsError_t cuvsCagraMerge(cuvsResources_t res, - cuvsCagraMergeParams_t params, + cuvsCagraIndexParams_t params, cuvsCagraIndex_t* indices, size_t num_indices, + cuvsFilter filter, cuvsCagraIndex_t output_index) { return cuvs::core::translate_exceptions([=] { @@ -621,16 +633,16 @@ extern "C" cuvsError_t cuvsCagraMerge(cuvsResources_t res, // Dispatch based on data type if (dtype.code == kDLFloat && dtype.bits == 32) { output_index->addr = - reinterpret_cast(_merge(res, *params, indices, num_indices)); + reinterpret_cast(_merge(res, *params, indices, num_indices, filter)); } else if (dtype.code == kDLFloat && dtype.bits == 16) { output_index->addr = - reinterpret_cast(_merge(res, *params, indices, num_indices)); + reinterpret_cast(_merge(res, *params, indices, num_indices, filter)); } else if (dtype.code == kDLInt && dtype.bits == 8) { output_index->addr = - reinterpret_cast(_merge(res, *params, indices, num_indices)); + reinterpret_cast(_merge(res, *params, indices, num_indices, filter)); } else if (dtype.code == kDLUInt && dtype.bits == 8) { output_index->addr = - reinterpret_cast(_merge(res, *params, indices, num_indices)); + reinterpret_cast(_merge(res, *params, indices, num_indices, filter)); } else { RAFT_FAIL("Unsupported index data type: code=%d, bits=%d", dtype.code, dtype.bits); } @@ -708,24 +720,6 @@ extern "C" cuvsError_t cuvsCagraSearchParamsDestroy(cuvsCagraSearchParams_t para return cuvs::core::translate_exceptions([=] { delete params; }); } -extern "C" cuvsError_t cuvsCagraMergeParamsCreate(cuvsCagraMergeParams_t* params) -{ - return cuvs::core::translate_exceptions([=] { - cuvsCagraIndexParams_t idx_params; - cuvsCagraIndexParamsCreate(&idx_params); - *params = new cuvsCagraMergeParams{.output_index_params = idx_params, - .strategy = MERGE_STRATEGY_PHYSICAL}; - }); -} - -extern "C" cuvsError_t cuvsCagraMergeParamsDestroy(cuvsCagraMergeParams_t params) -{ - return cuvs::core::translate_exceptions([=] { - cuvsCagraIndexParamsDestroy(params->output_index_params); - delete params; - }); -} - extern "C" cuvsError_t cuvsCagraDeserialize(cuvsResources_t res, const char* filename, cuvsCagraIndex_t index) diff --git a/c/tests/neighbors/ann_cagra_c.cu b/c/tests/neighbors/ann_cagra_c.cu index a39e0553ae..c1c74feb58 100644 --- a/c/tests/neighbors/ann_cagra_c.cu +++ b/c/tests/neighbors/ann_cagra_c.cu @@ -509,13 +509,15 @@ TEST(CagraC, BuildMergeSearch) ASSERT_EQ(cuvsCagraBuild(res, build_params, &main_dataset_tensor, index_main), CUVS_SUCCESS); ASSERT_EQ(cuvsCagraBuild(res, build_params, &additional_dataset_tensor, index_add), CUVS_SUCCESS); - cuvsCagraMergeParams_t merge_params; - cuvsCagraMergeParamsCreate(&merge_params); cuvsCagraIndex_t index_merged; cuvsCagraIndexCreate(&index_merged); + cuvsFilter filter; + filter.type = NO_FILTER; + filter.addr = 0; + cuvsCagraIndex_t index_array[2] = {index_main, index_add}; - ASSERT_EQ(cuvsCagraMerge(res, merge_params, index_array, 2, index_merged), CUVS_SUCCESS); + ASSERT_EQ(cuvsCagraMerge(res, build_params, index_array, 2, filter, index_merged), CUVS_SUCCESS); int64_t merged_dim = -1; ASSERT_EQ(cuvsCagraIndexGetDims(index_merged, &merged_dim), CUVS_SUCCESS); @@ -547,9 +549,6 @@ TEST(CagraC, BuildMergeSearch) cuvsCagraSearchParamsCreate(&search_params); (*search_params).itopk_size = 1; - cuvsFilter filter; - filter.type = NO_FILTER; - filter.addr = 0; ASSERT_EQ(cuvsCagraSearch(res, search_params, index_merged, @@ -569,7 +568,6 @@ TEST(CagraC, BuildMergeSearch) EXPECT_NEAR(distance_host, 0.0f, 1e-6); cuvsCagraSearchParamsDestroy(search_params); - cuvsCagraMergeParamsDestroy(merge_params); cuvsCagraIndexParamsDestroy(build_params); cuvsCagraIndexDestroy(index_merged); cuvsCagraIndexDestroy(index_add); diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index c8d1393046..177f995b17 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -230,38 +230,6 @@ struct extend_params { * 0. */ uint32_t max_chunk_size = 0; }; -/** - * @} - */ - -/** - * @defgroup cagra_cpp_merge_params CAGRA index merge parameters - * @{ - */ - -/** - * @brief Parameters for merging CAGRA indexes. - */ -struct merge_params : cuvs::neighbors::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::MERGE_STRATEGY_PHYSICAL`. - cuvs::neighbors::MergeStrategy merge_strategy = - cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL; - - /// Implementation of the polymorphic strategy() method - cuvs::neighbors::MergeStrategy strategy() const { return merge_strategy; } -}; - /** * @} */ @@ -2057,126 +2025,47 @@ void serialize_to_hnswlib( * 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); + * auto merged_index = cagra::merge(res, index_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. - * + * @param[in] row_filter an optional device filter function object that greenlights rows + * to include in the merged index (none_sample_filter for no filtering) * @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) + const cuvs::neighbors::cagra::index_params& params, + std::vector*>& indices, + const cuvs::neighbors::filtering::base_filter& row_filter = + cuvs::neighbors::filtering::none_sample_filter{}) -> 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. - */ +/** @copydoc merge */ auto merge(raft::resources const& res, - const cuvs::neighbors::cagra::merge_params& params, - std::vector*>& indices) + const cuvs::neighbors::cagra::index_params& params, + std::vector*>& indices, + const cuvs::neighbors::filtering::base_filter& row_filter = + cuvs::neighbors::filtering::none_sample_filter{}) -> 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. - */ +/** @copydoc merge */ auto merge(raft::resources const& res, - const cuvs::neighbors::cagra::merge_params& params, - std::vector*>& indices) + const cuvs::neighbors::cagra::index_params& params, + std::vector*>& indices, + const cuvs::neighbors::filtering::base_filter& row_filter = + cuvs::neighbors::filtering::none_sample_filter{}) -> 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. - */ +/** @copydoc merge */ auto merge(raft::resources const& res, - const cuvs::neighbors::cagra::merge_params& params, - std::vector*>& indices) + const cuvs::neighbors::cagra::index_params& params, + std::vector*>& indices, + const cuvs::neighbors::filtering::base_filter& row_filter = + cuvs::neighbors::filtering::none_sample_filter{}) -> cuvs::neighbors::cagra::index; /** * @} diff --git a/cpp/include/cuvs/neighbors/cagra_index_wrapper.hpp b/cpp/include/cuvs/neighbors/cagra_index_wrapper.hpp index c7223d12cf..f8eab6a684 100644 --- a/cpp/include/cuvs/neighbors/cagra_index_wrapper.hpp +++ b/cpp/include/cuvs/neighbors/cagra_index_wrapper.hpp @@ -23,11 +23,43 @@ namespace cuvs::neighbors::cagra { template struct index; -struct merge_params; +struct index_params; } // namespace cuvs::neighbors::cagra namespace cuvs::neighbors::cagra { +/** + * @defgroup cagra_cpp_merge_params CAGRA index merge parameters + * @{ + */ + +/** + * @brief Parameters for merging CAGRA indexes. + */ +struct merge_params : cuvs::neighbors::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::MERGE_STRATEGY_PHYSICAL`. + cuvs::neighbors::MergeStrategy merge_strategy = + cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL; + + /// Implementation of the polymorphic strategy() method + cuvs::neighbors::MergeStrategy strategy() const { return merge_strategy; } +}; + +/** + * @} + */ + /** * @brief Wrapper for CAGRA index implementing IndexWrapper. * diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index aaf1edb4e3..788240571a 100644 --- a/cpp/src/neighbors/cagra.cuh +++ b/cpp/src/neighbors/cagra.cuh @@ -408,12 +408,23 @@ void extend( template index merge(raft::resources const& handle, - const cagra::merge_params& params, - std::vector*>& indices) + const cagra::index_params& params, + std::vector*>& indices, + const cuvs::neighbors::filtering::base_filter& row_filter) { - return cagra::detail::merge(handle, params, indices); + return cagra::detail::merge(handle, params, indices, row_filter); } /** @} */ // end group cagra } // namespace cuvs::neighbors::cagra + +#define CUVS_INST_CAGRA_MERGE(T, IdxT) \ + auto merge(raft::resources const& handle, \ + const cuvs::neighbors::cagra::index_params& params, \ + std::vector*>& indices, \ + const cuvs::neighbors::filtering::base_filter& row_filter) \ + -> cuvs::neighbors::cagra::index \ + { \ + return cuvs::neighbors::cagra::merge(handle, params, indices, row_filter); \ + } diff --git a/cpp/src/neighbors/cagra_index_wrapper.cu b/cpp/src/neighbors/cagra_index_wrapper.cu index 445ed77cec..39926f9a22 100644 --- a/cpp/src/neighbors/cagra_index_wrapper.cu +++ b/cpp/src/neighbors/cagra_index_wrapper.cu @@ -89,8 +89,9 @@ IndexWrapper::merge( return std::make_shared>( std::move(wrappers)); } else if (cagra_params->strategy() == cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL) { - auto merged_index = cuvs::neighbors::cagra::merge(handle, *cagra_params, cagra_indices); - auto* idx = new decltype(merged_index)(std::move(merged_index)); + auto merged_index = + cuvs::neighbors::cagra::merge(handle, cagra_params->output_index_params, cagra_indices); + auto* idx = new decltype(merged_index)(std::move(merged_index)); return std::make_shared>(idx); } diff --git a/cpp/src/neighbors/cagra_merge_float.cu b/cpp/src/neighbors/cagra_merge_float.cu index 01526d7b3c..e7e3a25775 100644 --- a/cpp/src/neighbors/cagra_merge_float.cu +++ b/cpp/src/neighbors/cagra_merge_float.cu @@ -18,18 +18,5 @@ #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 - +CUVS_INST_CAGRA_MERGE(float, uint32_t); } // namespace cuvs::neighbors::cagra diff --git a/cpp/src/neighbors/cagra_merge_half.cu b/cpp/src/neighbors/cagra_merge_half.cu index 90bdafb255..8df3b2ccf2 100644 --- a/cpp/src/neighbors/cagra_merge_half.cu +++ b/cpp/src/neighbors/cagra_merge_half.cu @@ -19,17 +19,6 @@ 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 +CUVS_INST_CAGRA_MERGE(half, uint32_t); } // namespace cuvs::neighbors::cagra diff --git a/cpp/src/neighbors/cagra_merge_int8.cu b/cpp/src/neighbors/cagra_merge_int8.cu index 854261336e..544c04dc8b 100644 --- a/cpp/src/neighbors/cagra_merge_int8.cu +++ b/cpp/src/neighbors/cagra_merge_int8.cu @@ -19,17 +19,6 @@ 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 +CUVS_INST_CAGRA_MERGE(int8_t, uint32_t); } // namespace cuvs::neighbors::cagra diff --git a/cpp/src/neighbors/cagra_merge_uint8.cu b/cpp/src/neighbors/cagra_merge_uint8.cu index d536029cc1..e42c4c0fe2 100644 --- a/cpp/src/neighbors/cagra_merge_uint8.cu +++ b/cpp/src/neighbors/cagra_merge_uint8.cu @@ -19,17 +19,6 @@ 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 +CUVS_INST_CAGRA_MERGE(uint8_t, uint32_t); } // namespace cuvs::neighbors::cagra diff --git a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh index 86faad02c0..03e6be3ca7 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh @@ -25,6 +25,7 @@ #include #include #include +#include #include #include @@ -41,13 +42,10 @@ namespace cuvs::neighbors::cagra::detail { template index merge(raft::resources const& handle, - const cagra::merge_params& params, - std::vector*>& indices) + const cagra::index_params& params, + std::vector*>& indices, + const cuvs::neighbors::filtering::base_filter& row_filter) { - // we're doing a physical merge here, make sure that this matches the merge_params - RAFT_EXPECTS(params.merge_strategy == cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL, - "cagra::merge only supports merge_strategy=MERGE_STRATEGY_PHYSICAL"); - using cagra_index_t = cuvs::neighbors::cagra::index; using ds_idx_type = typename cagra_index_t::dataset_index_type; @@ -55,6 +53,9 @@ index merge(raft::resources const& handle, std::size_t new_dataset_size = 0; int64_t stride = -1; + RAFT_EXPECTS(row_filter.get_filter_type() != cuvs::neighbors::filtering::FilterType::Bitmap, + "Bitmap filter isn't supported inside cagra::merge"); + for (cagra_index_t* index : indices) { RAFT_EXPECTS(index != nullptr, "Null pointer detected in 'indices'. Ensure all elements are valid before usage."); @@ -95,29 +96,74 @@ index merge(raft::resources const& handle, } }; - 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)); + auto updated_dataset = + raft::make_device_matrix(handle, int64_t(new_dataset_size), 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}); + if (row_filter.get_filter_type() == cuvs::neighbors::filtering::FilterType::Bitset) { + auto actual_filter = + dynamic_cast&>( + row_filter); + auto filtered_row_count = actual_filter.view().count(handle); + + // Convert the filter to a CSR matrix (so that we can pass indices to raft::copy_rows) + auto indices_csr = raft::make_device_csr_matrix( + handle, 1, new_dataset_size); + indices_csr.initialize_sparsity(filtered_row_count); + + actual_filter.view().to_csr(handle, indices_csr); + + // Get the indices array from the csr matrix. Note that this returns a raft::span object + // and we need to pass as device_vector_view, which is a 1D mdspan (instead of a span) + // so we need to translate here (and adjust to be const) + auto indices = indices_csr.structure_view().get_indices(); + auto indices_view = raft::make_device_vector_view( + indices.data(), static_cast(indices.size())); + + auto filtered_dataset = raft::make_device_matrix(handle, filtered_row_count, dim); + raft::matrix::copy_rows(handle, + raft::make_const_mdspan(updated_dataset.view()), + filtered_dataset.view(), + indices_view); + + auto merged_index = + cagra::build(handle, params, raft::make_const_mdspan(filtered_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(filtered_dataset.view().extents(), + std::array{stride, 1}); + + merged_index.update_dataset(handle, owning_t{std::move(filtered_dataset), out_layout}); + } + RAFT_LOG_DEBUG("cagra merge: using device memory for merged dataset"); + return merged_index; + } else { + 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}); + } + RAFT_LOG_DEBUG("cagra merge: using device memory for merged dataset"); + return merged_index; } - RAFT_LOG_DEBUG("cagra merge: using device memory for merged dataset"); - return merged_index; - } catch (std::bad_alloc& e) { + // We don't currently support the cpu memory fallback with filtered merge, since the + // 'raft::matrix::copy_rows' only supports gpu memory + RAFT_EXPECTS(row_filter.get_filter_type() == cuvs::neighbors::filtering::FilterType::None, + "Filtered merge isn't available on cpu memory"); + RAFT_LOG_DEBUG("cagra::merge: using host memory for merged dataset"); auto updated_dataset = @@ -126,8 +172,8 @@ index merge(raft::resources const& handle, 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) { + 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; diff --git a/cpp/tests/neighbors/ann_cagra.cuh b/cpp/tests/neighbors/ann_cagra.cuh index 02a090ea8e..67929b98f1 100644 --- a/cpp/tests/neighbors/ann_cagra.cuh +++ b/cpp/tests/neighbors/ann_cagra.cuh @@ -943,6 +943,242 @@ class AnnCagraFilterTest : public ::testing::TestWithParam { rmm::device_uvector search_queries; }; +template +class AnnCagraIndexFilteredMergeTest : public ::testing::TestWithParam { + public: + AnnCagraIndexFilteredMergeTest() + : stream_(raft::resource::get_cuda_stream(handle_)), + ps(::testing::TestWithParam::GetParam()), + database(0, stream_), + search_queries(0, stream_) + { + } + + protected: + template + void testCagra() + { + if (ps.metric == cuvs::distance::DistanceType::CosineExpanded) { + if (ps.build_algo == graph_build_algo::ITERATIVE_CAGRA_SEARCH || ps.dim == 1) { + GTEST_SKIP(); + } + } + if (ps.compression != std::nullopt) GTEST_SKIP(); + // IVF_PQ graph build does not support BitwiseHamming + if (ps.metric == cuvs::distance::DistanceType::BitwiseHamming && + ((!std::is_same_v) || (ps.build_algo == graph_build_algo::IVF_PQ))) + 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(); + + // Avoid splitting datasets with a size of 0 + if (ps.n_rows <= 3) GTEST_SKIP(); + + // IVF_PQ requires the `n_rows >= n_lists`. + if (ps.n_rows < 8 && ps.build_algo == graph_build_algo::IVF_PQ) GTEST_SKIP(); + + // can only use physical merge for filtered merge + if (ps.merge_strategy != cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL) { + 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); + + // Create a bitset filter to test out the merge + auto removed_indices = + raft::make_device_vector(handle_, test_cagra_sample_filter::offset); + thrust::sequence( + raft::resource::get_thrust_policy(handle_), + thrust::device_pointer_cast(removed_indices.data_handle()), + thrust::device_pointer_cast(removed_indices.data_handle() + removed_indices.extent(0))); + raft::resource::sync_stream(handle_); + cuvs::core::bitset removed_indices_bitset( + handle_, removed_indices.view(), ps.n_rows); + auto bitset_filter_obj = + cuvs::neighbors::filtering::bitset_filter(removed_indices_bitset.view()); + + { + rmm::device_uvector distances_naive_dev(queries_size, stream_); + rmm::device_uvector indices_naive_dev(queries_size, stream_); + + auto* database_filtered_ptr = database.data() + test_cagra_sample_filter::offset * ps.dim; + cuvs::neighbors::naive_knn( + handle_, + distances_naive_dev.data(), + indices_naive_dev.data(), + search_queries.data(), + database_filtered_ptr, + ps.n_queries, + ps.n_rows - test_cagra_sample_filter::offset, + ps.dim, + ps.k, + ps.metric); + raft::linalg::addScalar(indices_naive_dev.data(), + indices_naive_dev.data(), + IdxT(test_cagra_sample_filter::offset), + queries_size, + stream_); + + 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), index_params.metric); + 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 split_ratio = 0.55; + const std::size_t database0_size = ps.n_rows * split_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_, index_params.metric); + cagra::index index1(handle_, index_params.metric); + std::optional> database_host{std::nullopt}; + if (ps.host_dataset) { + database_host = raft::make_host_matrix(handle_, ps.n_rows, ps.dim); + raft::copy(database_host->data_handle(), database.data(), database.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); + } + { + auto database_host_view = raft::make_host_matrix_view( + (const DataT*)database_host->data_handle() + database0_size * ps.dim, + 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; + indices.push_back(&index0); + indices.push_back(&index1); + + auto index = + cuvs::neighbors::cagra::merge(handle_, index_params, indices, bitset_filter_obj); + + 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; + + cuvs::neighbors::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)); + + /* TODO: eval_distances doesn't work, potentially because of id translation mismatch + if (!ps.compression.has_value()) { + 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; +}; + template class AnnCagraIndexMergeTest : public ::testing::TestWithParam { public: @@ -1039,8 +1275,8 @@ class AnnCagraIndexMergeTest : public ::testing::TestWithParam { break; }; - const double splite_ratio = 0.55; - const std::size_t database0_size = ps.n_rows * splite_ratio; + const double split_ratio = 0.55; + const std::size_t database0_size = ps.n_rows * split_ratio; const std::size_t database1_size = ps.n_rows - database0_size; auto database0_view = raft::make_device_matrix_view( 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 501550c3b0..98ce6c9163 100644 --- a/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu +++ b/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu @@ -1,5 +1,5 @@ /* - * Copyright (c) 2023-2024, NVIDIA CORPORATION. + * Copyright (c) 2023-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. @@ -34,6 +34,13 @@ typedef AnnCagraIndexMergeTest AnnCagraIndexMergeTe TEST_P(AnnCagraIndexMergeTestF_U32, AnnCagraIndexMerge_U32) { this->testCagra(); } TEST_P(AnnCagraIndexMergeTestF_U32, AnnCagraIndexMerge_I64) { this->testCagra(); } +typedef AnnCagraIndexFilteredMergeTest + AnnCagraIndexFilteredMergeTestF_U32; +TEST_P(AnnCagraIndexFilteredMergeTestF_U32, AnnCagraIndexFilteredMerge_U32) +{ + this->testCagra(); +} + INSTANTIATE_TEST_CASE_P(AnnCagraTest, AnnCagraTestF_U32, ::testing::ValuesIn(inputs)); INSTANTIATE_TEST_CASE_P(AnnCagraAddNodesTest, AnnCagraAddNodesTestF_U32, @@ -45,4 +52,8 @@ INSTANTIATE_TEST_CASE_P(AnnCagraIndexMergeTest, AnnCagraIndexMergeTestF_U32, ::testing::ValuesIn(inputs)); +INSTANTIATE_TEST_CASE_P(AnnCagraIndexFilteredMergeTest, + AnnCagraIndexFilteredMergeTestF_U32, + ::testing::ValuesIn(inputs)); + } // namespace cuvs::neighbors::cagra diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java index a3cb68448c..65320e24d1 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/CagraIndex.java @@ -187,7 +187,7 @@ static CagraIndex merge(CagraIndex[] indexes) throws Throwable { * @return A new merged CAGRA index * @throws Throwable if an error occurs during the merge operation */ - static CagraIndex merge(CagraIndex[] indexes, CagraMergeParams mergeParams) throws Throwable { + static CagraIndex merge(CagraIndex[] indexes, CagraIndexParams mergeParams) throws Throwable { if (indexes == null || indexes.length == 0) { throw new IllegalArgumentException("At least one index must be provided for merging"); } diff --git a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java index 034aaf0d7d..2bbc1543b2 100644 --- a/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java +++ b/java/cuvs-java/src/main/java/com/nvidia/cuvs/spi/CuVSProvider.java @@ -145,7 +145,7 @@ TieredIndex.Builder newTieredIndexBuilder(CuVSResources cuVSResources) * @return A new merged CAGRA index * @throws Throwable if an error occurs during the merge operation */ - default CagraIndex mergeCagraIndexes(CagraIndex[] indexes, CagraMergeParams mergeParams) + default CagraIndex mergeCagraIndexes(CagraIndex[] indexes, CagraIndexParams mergeParams) throws Throwable { // Default implementation falls back to the method without parameters return mergeCagraIndexes(indexes); diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSParamsHelper.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSParamsHelper.java index 6774cc4619..87dc463c18 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSParamsHelper.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/internal/CuVSParamsHelper.java @@ -114,25 +114,6 @@ public void close() { } } - static CloseableHandle createCagraMergeParams() { - try (var localArena = Arena.ofConfined()) { - var paramsPtrPtr = localArena.allocate(cuvsCagraMergeParams_t); - checkCuVSError(cuvsCagraMergeParamsCreate(paramsPtrPtr), "cuvsCagraMergeParamsCreate"); - var paramsPtr = paramsPtrPtr.get(cuvsCagraMergeParams_t, 0L); - return new CloseableHandle() { - @Override - public MemorySegment handle() { - return paramsPtr; - } - - @Override - public void close() { - checkCuVSError(cuvsCagraMergeParamsDestroy(paramsPtr), "cuvsCagraMergeParamsDestroy"); - } - }; - } - } - static CloseableHandle createHnswIndexParams() { try (var localArena = Arena.ofConfined()) { var paramsPtrPtr = localArena.allocate(cuvsHnswIndexParams_t); diff --git a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java index 8a036a62ef..44351419fa 100644 --- a/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java +++ b/java/cuvs-java/src/main/java22/com/nvidia/cuvs/spi/JDKProvider.java @@ -174,7 +174,7 @@ public CagraIndex mergeCagraIndexes(CagraIndex[] indexes) { } @Override - public CagraIndex mergeCagraIndexes(CagraIndex[] indexes, CagraMergeParams mergeParams) { + public CagraIndex mergeCagraIndexes(CagraIndex[] indexes, CagraIndexParams mergeParams) { if (indexes == null || indexes.length == 0) { throw new IllegalArgumentException("At least one index must be provided for merging"); }