diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt index a4a60d6a7f..52858db577 100644 --- a/cpp/CMakeLists.txt +++ b/cpp/CMakeLists.txt @@ -419,6 +419,9 @@ if(BUILD_SHARED_LIBS) src/neighbors/detail/cagra/cagra_build.cpp src/neighbors/detail/cagra/topk_for_cagra/topk.cu src/neighbors/dynamic_batching.cu + src/neighbors/cagra_index_wrapper.cu + src/neighbors/composite/index.cu + src/neighbors/composite/merge.cpp $<$:src/neighbors/hnsw.cpp> src/neighbors/ivf_flat_index.cpp src/neighbors/ivf_flat/ivf_flat_build_extend_float_int64_t.cu diff --git a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h index 48a0f39ab7..7ad8ef91b5 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h @@ -277,6 +277,20 @@ void parse_build_param(const nlohmann::json& conf, parse_build_param(comp_search_conf, vpq_pams); param.cagra_params.compression.emplace(vpq_pams); } + + if (conf.contains("num_dataset_splits")) { + param.num_dataset_splits = conf.at("num_dataset_splits"); + } + if (conf.contains("merge_type")) { + std::string mt = conf.at("merge_type"); + if (mt == "PHYSICAL") { + param.merge_type = cuvs::bench::CagraMergeType::kPhysical; + } else if (mt == "LOGICAL") { + param.merge_type = cuvs::bench::CagraMergeType::kLogical; + } else { + throw std::runtime_error("invalid value for merge_type"); + } + } } cuvs::bench::AllocatorType parse_allocator(std::string mem_type) diff --git a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h index e4c441639c..ed21179b5e 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h @@ -24,6 +24,7 @@ #include #include #include +#include #include #include #include @@ -44,14 +45,17 @@ #include #include #include +#include #include #include #include +#include namespace cuvs::bench { enum class AllocatorType { kHostPinned, kHostHugePage, kDevice }; enum class CagraBuildAlgo { kAuto, kIvfPq, kNnDescent }; +enum class CagraMergeType { kPhysical, kLogical }; template class cuvs_cagra : public algo, public algo_gpu { @@ -80,6 +84,8 @@ class cuvs_cagra : public algo, public algo_gpu { std::optional ivf_pq_refine_rate = std::nullopt; std::optional ivf_pq_build_params = std::nullopt; std::optional ivf_pq_search_params = std::nullopt; + size_t num_dataset_splits = 1; + CagraMergeType merge_type = CagraMergeType::kPhysical; void prepare_build_params(const raft::extent_2d& dataset_extents) { @@ -188,6 +194,7 @@ class cuvs_cagra : public algo, public algo_gpu { bool dynamic_batching_conservative_dispatch_; std::shared_ptr filter_; + std::vector>> sub_indices_; inline rmm::device_async_resource_ref get_mr(AllocatorType mem_type) { @@ -211,10 +218,57 @@ void cuvs_cagra::build(const T* dataset, size_t nrow) auto dataset_view_device = raft::make_mdspan(dataset, dataset_extents); bool dataset_is_on_host = raft::get_device_for_address(dataset) == -1; + if (index_params_.num_dataset_splits <= 1) { + index_ = std::make_shared>(std::move( + dataset_is_on_host ? cuvs::neighbors::cagra::build(handle_, params, dataset_view_host) + : cuvs::neighbors::cagra::build(handle_, params, dataset_view_device))); + } else { + IdxT rows_per_split = + raft::ceildiv(nrow, static_cast(index_params_.num_dataset_splits)); + for (size_t i = 0; i < index_params_.num_dataset_splits; ++i) { + IdxT start = static_cast(i * rows_per_split); + if (start >= nrow) break; + IdxT rows = std::min(rows_per_split, static_cast(nrow) - start); + const T* sub_ptr = dataset + static_cast(start) * dimension_; + auto sub_host = + raft::make_host_matrix_view(sub_ptr, rows, dimension_); + auto sub_dev = + raft::make_device_matrix_view(sub_ptr, rows, dimension_); + + auto sub_index = + cuvs::neighbors::cagra::index(handle_, index_params_.cagra_params.metric); + if (index_params_.merge_type == CagraMergeType::kPhysical) { + if (dataset_is_on_host) { + sub_index.update_dataset(handle_, sub_host); + } else { + sub_index.update_dataset(handle_, sub_dev); + } + } + if (index_params_.merge_type == CagraMergeType::kLogical) { + if (dataset_is_on_host) { + sub_index = cuvs::neighbors::cagra::build(handle_, params, sub_host); + } else { + sub_index = cuvs::neighbors::cagra::build(handle_, params, sub_dev); + } + } + auto sub_index_shared = + std::make_shared>(std::move(sub_index)); + sub_indices_.push_back(std::move(sub_index_shared)); + } + if (index_params_.merge_type == CagraMergeType::kPhysical) { + cuvs::neighbors::cagra::merge_params merge_params{index_params_.cagra_params}; + merge_params.merge_strategy = cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL; + + std::vector*> indices; + indices.reserve(sub_indices_.size()); + for (auto& ptr : sub_indices_) { + indices.push_back(ptr.get()); + } - index_ = std::make_shared>(std::move( - dataset_is_on_host ? cuvs::neighbors::cagra::build(handle_, params, dataset_view_host) - : cuvs::neighbors::cagra::build(handle_, params, dataset_view_device))); + index_ = std::make_shared>( + std::move(cuvs::neighbors::cagra::merge(handle_, merge_params, indices))); + } + } } inline auto allocator_to_string(AllocatorType mem_type) -> std::string @@ -233,7 +287,7 @@ template void cuvs_cagra::set_search_param(const search_param_base& param, const void* filter_bitset) { - filter_ = make_cuvs_filter(filter_bitset, index_->size()); + if (index_) { filter_ = make_cuvs_filter(filter_bitset, index_->size()); } auto sp = dynamic_cast(param); bool needs_dynamic_batcher_update = (dynamic_batching_max_batch_size_ != sp.dynamic_batching_max_batch_size) || @@ -314,27 +368,65 @@ void cuvs_cagra::set_search_param(const search_param_base& param, template void cuvs_cagra::set_search_dataset(const T* dataset, size_t nrow) { - using ds_idx_type = decltype(index_->data().n_rows()); - bool is_vpq = - dynamic_cast*>(&index_->data()) || - dynamic_cast*>(&index_->data()); - // It can happen that we are re-using a previous algo object which already has - // the dataset set. Check if we need update. - if (static_cast(input_dataset_v_->extent(0)) != nrow || - input_dataset_v_->data_handle() != dataset) { - *input_dataset_v_ = raft::make_device_matrix_view(dataset, nrow, this->dim_); - need_dataset_update_ = !is_vpq; // ignore update if this is a VPQ dataset. + if (index_params_.num_dataset_splits > 1 && + index_params_.merge_type == CagraMergeType::kLogical) { + bool dataset_is_on_host = raft::get_device_for_address(dataset) == -1; + IdxT rows_per_split = + raft::ceildiv(nrow, static_cast(index_params_.num_dataset_splits)); + for (size_t i = 0; i < sub_indices_.size(); ++i) { + IdxT start = static_cast(i * rows_per_split); + if (start >= nrow) break; + IdxT rows = std::min(rows_per_split, static_cast(nrow) - start); + const T* sub_ptr = dataset + static_cast(start) * dimension_; + auto sub_host = + raft::make_host_matrix_view(sub_ptr, rows, dimension_); + auto sub_dev = + raft::make_device_matrix_view(sub_ptr, rows, dimension_); + auto sub_index = sub_indices_[i].get(); + if (index_params_.merge_type == CagraMergeType::kLogical) { + if (dataset_is_on_host) { + sub_index->update_dataset(handle_, sub_host); + } else { + sub_index->update_dataset(handle_, sub_dev); + } + } + } + need_dataset_update_ = false; + } else { + using ds_idx_type = decltype(index_->data().n_rows()); + bool is_vpq = + dynamic_cast*>(&index_->data()) || + dynamic_cast*>(&index_->data()); + // It can happen that we are re-using a previous algo object which already has + // the dataset set. Check if we need update. + if (static_cast(input_dataset_v_->extent(0)) != nrow || + input_dataset_v_->data_handle() != dataset) { + *input_dataset_v_ = + raft::make_device_matrix_view(dataset, nrow, this->dim_); + need_dataset_update_ = !is_vpq; // ignore update if this is a VPQ dataset. + } } } template void cuvs_cagra::save(const std::string& file) const { - using ds_idx_type = decltype(index_->data().n_rows()); - bool is_vpq = - dynamic_cast*>(&index_->data()) || - dynamic_cast*>(&index_->data()); - cuvs::neighbors::cagra::serialize(handle_, file, *index_, is_vpq); + if (index_params_.num_dataset_splits > 1 && + index_params_.merge_type == CagraMergeType::kLogical) { + for (size_t i = 0; i < sub_indices_.size(); ++i) { + std::string subfile = file + (i == 0 ? "" : ".subidx." + std::to_string(i)); + cuvs::neighbors::cagra::serialize(handle_, subfile, *sub_indices_[i], false); + } + std::ofstream f(file + ".submeta", std::ios::out); + f << sub_indices_.size(); + f.close(); + } else { + using ds_idx_type = decltype(index_->data().n_rows()); + bool is_vpq = + dynamic_cast*>(&index_->data()) || + dynamic_cast*>(&index_->data()); + cuvs::neighbors::cagra::serialize(handle_, file, *index_, is_vpq); + } } template @@ -346,8 +438,24 @@ void cuvs_cagra::save_to_hnswlib(const std::string& file) const template void cuvs_cagra::load(const std::string& file) { - index_ = std::make_shared>(handle_); - cuvs::neighbors::cagra::deserialize(handle_, file, index_.get()); + std::ifstream meta(file + ".submeta", std::ios::in); + if (index_params_.num_dataset_splits > 1 && + index_params_.merge_type == CagraMergeType::kLogical && meta.good()) { + // Load multiple sub-indices for logical merge + size_t count; + meta >> count; + meta.close(); + sub_indices_.clear(); + for (size_t i = 0; i < count; ++i) { + std::string subfile = file + (i == 0 ? "" : ".subidx." + std::to_string(i)); + auto sub_index = std::make_shared>(handle_); + cuvs::neighbors::cagra::deserialize(handle_, subfile, sub_index.get()); + sub_indices_.push_back(std::move(sub_index)); + } + } else { + index_ = std::make_shared>(handle_); + cuvs::neighbors::cagra::deserialize(handle_, file, index_.get()); + } } template @@ -377,8 +485,41 @@ void cuvs_cagra::search_base( neighbors_view, distances_view); } else { - cuvs::neighbors::cagra::search( - handle_, search_params_, *index_, queries_view, neighbors_view, distances_view, *filter_); + if (index_params_.num_dataset_splits <= 1 || + index_params_.merge_type == CagraMergeType::kPhysical) { + cuvs::neighbors::cagra::search( + handle_, search_params_, *index_, queries_view, neighbors_view, distances_view, *filter_); + } else { + if (index_params_.merge_type == CagraMergeType::kLogical) { + cuvs::neighbors::cagra::merge_params merge_params{index_params_.cagra_params}; + merge_params.merge_strategy = cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_LOGICAL; + + // Create wrapped indices for composite merge + std::vector>> + wrapped_indices; + wrapped_indices.reserve(sub_indices_.size()); + for (auto& ptr : sub_indices_) { + auto index_wrapper = + cuvs::neighbors::cagra::make_index_wrapper(ptr.get()); + wrapped_indices.push_back(index_wrapper); + } + + raft::resources composite_handle(handle_); + size_t n_streams = wrapped_indices.size(); + raft::resource::set_cuda_stream_pool(composite_handle, + std::make_shared(n_streams)); + + auto merged_index = + cuvs::neighbors::composite::merge(composite_handle, merge_params, wrapped_indices); + cuvs::neighbors::filtering::none_sample_filter empty_filter; + merged_index->search(composite_handle, + search_params_, + queries_view, + neighbors_view, + distances_view, + empty_filter); + } + } } } diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index 036dbba9c8..488c769312 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -18,7 +18,6 @@ #include "common.hpp" #include -#include #include #include #include @@ -275,17 +274,10 @@ struct extend_params { * @{ */ -/** - * @brief Determines the strategy for merging CAGRA graphs. - * - * @note Currently, only the MERGE_STRATEGY_PHYSICAL strategy is supported. - */ -using MergeStrategy = cuvsMergeStrategy; - /** * @brief Parameters for merging CAGRA indexes. */ -struct merge_params { +struct merge_params : cuvs::neighbors::merge_params { merge_params() = default; /** @@ -298,7 +290,11 @@ struct merge_params { cagra::index_params output_index_params; /// Strategy for merging. Defaults to `MergeStrategy::MERGE_STRATEGY_PHYSICAL`. - MergeStrategy strategy = 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; } }; /** @@ -2602,3 +2598,5 @@ auto distribute(const raft::resources& clique, const std::string& filename) -> cuvs::neighbors::mg_index, T, IdxT>; } // namespace cuvs::neighbors::cagra + +#include diff --git a/cpp/include/cuvs/neighbors/cagra_index_wrapper.hpp b/cpp/include/cuvs/neighbors/cagra_index_wrapper.hpp new file mode 100644 index 0000000000..c7223d12cf --- /dev/null +++ b/cpp/include/cuvs/neighbors/cagra_index_wrapper.hpp @@ -0,0 +1,142 @@ +/* + * 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. + */ + +#pragma once + +#include +#include + +// Forward declarations to avoid circular dependencies +namespace cuvs::neighbors::cagra { +template +struct index; +struct merge_params; +} // namespace cuvs::neighbors::cagra + +namespace cuvs::neighbors::cagra { + +/** + * @brief Wrapper for CAGRA index implementing IndexWrapper. + * + * This class wraps a CAGRA index and provides compatibility with the IndexBase interface. + * It serves as a bridge to help the CAGRA index implementation transition from its + * original design to the new object-oriented polymorphic design based on IndexBase. + * + * The wrapper enables: + * - CAGRA index to work seamlessly with the new polymorphic IndexBase interface + * - Gradual migration from the original CAGRA API to the unified index architecture + * - Compatibility with composite index patterns and other polymorphic usage scenarios + * - Preservation of existing CAGRA functionality while adopting the new design patterns + * + * This allows existing CAGRA users to benefit from the new architecture without + * requiring immediate changes to their existing code, while new users can adopt + * the unified interface from the start. + */ +template +class IndexWrapper : public cuvs::neighbors::IndexWrapper { + public: + using base_type = cuvs::neighbors::IndexWrapper; + using value_type = typename base_type::value_type; + using index_type = typename base_type::index_type; + using out_index_type = typename base_type::out_index_type; + using matrix_index_type = typename base_type::matrix_index_type; + + explicit IndexWrapper(cuvs::neighbors::cagra::index* idx); + + void search( + const raft::resources& handle, + const cuvs::neighbors::search_params& params, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& filter = + cuvs::neighbors::filtering::none_sample_filter{}) const override; + + index_type size() const noexcept override; + + cuvs::distance::DistanceType metric() const noexcept override; + + /** + * @brief Merge this CAGRA index with other CAGRA indices. + * + * This method provides merge capability for CAGRA indices. It supports both + * physical merge (calling native CAGRA merge) and logical merge (creating + * CompositeIndex with wrapped indices). + * + * @param[in] handle RAFT resources for executing operations + * @param[in] params Merge parameters containing strategy and CAGRA-specific settings + * @param[in] other_indices Vector of other indices to merge with this one + * @return Shared pointer to merged index + */ + std::shared_ptr> merge( + const raft::resources& handle, + const cuvs::neighbors::merge_params& params, + const std::vector< + std::shared_ptr>>& + other_indices) const override; + + protected: + const cuvs::neighbors::search_params& convert_search_params( + const cuvs::neighbors::search_params& params) const override + { + // For CAGRA, we expect the params to be cagra::search_params + // This is handled in the search method via static_cast + return params; + } + + private: + cuvs::neighbors::cagra::index* index_; +}; + +/** + * @brief Factory function for creating a wrapped CAGRA index. + * + * This function creates a shared pointer to an IndexWrapper that wraps a CAGRA index, + * enabling it to work with the polymorphic IndexBase interface and composite operations. + * + * @tparam T Data type + * @tparam IdxT Index type + * @tparam OutputIdxT Output index type + * @param index Pointer to the CAGRA index + * @return Shared pointer to the wrapped index + * + * @par Example usage: + * @code{.cpp} + * // Create multiple CAGRA indices + * auto cagra_index1 = cuvs::neighbors::cagra::build(res, params, dataset1); + * auto cagra_index2 = cuvs::neighbors::cagra::build(res, params, dataset2); + * + * // Wrap them for polymorphic usage + * auto wrapped_index1 = cuvs::neighbors::cagra::make_index_wrapper(&cagra_index1); + * auto wrapped_index2 = cuvs::neighbors::cagra::make_index_wrapper(&cagra_index2); + * + * // Merge indices using the composite merge function + * std::vector>> indices; + * indices.push_back(wrapped_index1); + * indices.push_back(wrapped_index2); + * + * cuvs::neighbors::cagra::merge_params merge_params; + * auto merged_index = cuvs::neighbors::composite::merge(res, merge_params, indices); + * @endcode + */ +template +inline auto make_index_wrapper(cuvs::neighbors::cagra::index* index) + -> std::shared_ptr> +{ + return std::make_shared>(index); +} + +} // namespace cuvs::neighbors::cagra diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index cb56e9f8c2..47155ed47e 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2024, NVIDIA CORPORATION. + * Copyright (c) 2024-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. @@ -96,6 +96,26 @@ struct index_params { struct search_params {}; +/** + * @brief Strategy for merging indices. + * + * This enum is declared separately to avoid namespace pollution when including common.hpp. + * It provides a generic merge strategy that can be used across different index types. + */ +enum class MergeStrategy { + /** Merge indices physically by combining their data structures */ + MERGE_STRATEGY_PHYSICAL = 0, + /** Merge indices logically by creating a composite wrapper */ + MERGE_STRATEGY_LOGICAL = 1 +}; + +/** Base merge parameters with polymorphic interface. */ +struct merge_params { + virtual ~merge_params() = default; + + virtual MergeStrategy strategy() const = 0; +}; + /** @} */ // end group neighbors_index /** Two-dimensional dataset; maybe owning, maybe compressed, maybe strided. */ diff --git a/cpp/include/cuvs/neighbors/composite/index.hpp b/cpp/include/cuvs/neighbors/composite/index.hpp new file mode 100644 index 0000000000..df256e43d0 --- /dev/null +++ b/cpp/include/cuvs/neighbors/composite/index.hpp @@ -0,0 +1,103 @@ +/* + * 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. + */ + +#pragma once + +#include +#include +#include + +#include +#include + +namespace cuvs::neighbors::composite { + +/** + * @brief Composite index made of other IndexBase implementations. + */ +template +class CompositeIndex : public IndexBase { + public: + using value_type = typename IndexBase::value_type; + using index_type = typename IndexBase::index_type; + using out_index_type = typename IndexBase::out_index_type; + using matrix_index_type = typename IndexBase::matrix_index_type; + + using index_ptr = std::shared_ptr>; + + explicit CompositeIndex(std::vector children) : children_(std::move(children)) {} + + /** + * @brief Search the composite index for the k nearest neighbors. + * + * When the composite index contains multiple sub-indices, the user can set a + * stream pool in the input raft::resource to enable parallel search across + * sub-indices for improved performance. + * + * Usage example: + * @code{.cpp} + * using namespace cuvs::neighbors; + * // create a composite index with multiple sub-indices + * std::vector::index_ptr> sub_indices; + * // ... populate sub_indices ... + * auto composite_index = CompositeIndex(std::move(sub_indices)); + * + * // optional: create a stream pool to enable parallel search across sub-indices + * // recommended stream count: min(number_of_sub_indices, 8) + * size_t n_streams = std::min(sub_indices.size(), size_t(8)); + * raft::resource::set_cuda_stream_pool(handle, + * std::make_shared(n_streams)); + * + * // perform search with parallel sub-index execution + * composite_index.search(handle, search_params, queries, neighbors, distances); + * @endcode + * + * @param[in] handle raft resource handle + * @param[in] params search parameters + * @param[in] queries device matrix view of query vectors [n_queries, dim] + * @param[out] neighbors device matrix view for neighbor indices [n_queries, k] + * @param[out] distances device matrix view for distances [n_queries, k] + * @param[in] filter optional filter for search results + */ + void search( + const raft::resources& handle, + const cuvs::neighbors::search_params& params, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& filter = + cuvs::neighbors::filtering::none_sample_filter{}) const override; + + index_type size() const noexcept override + { + index_type total = 0; + for (const auto& c : children_) { + total += c->size(); + } + return total; + } + + cuvs::distance::DistanceType metric() const noexcept override + { + return children_.empty() ? cuvs::distance::DistanceType::L2Expanded + : children_.front()->metric(); + } + + private: + std::vector children_; +}; + +} // namespace cuvs::neighbors::composite diff --git a/cpp/include/cuvs/neighbors/composite/merge.hpp b/cpp/include/cuvs/neighbors/composite/merge.hpp new file mode 100644 index 0000000000..6635b7722e --- /dev/null +++ b/cpp/include/cuvs/neighbors/composite/merge.hpp @@ -0,0 +1,49 @@ +/* + * 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. + */ + +#pragma once + +#include "../index_base.hpp" +#include "../index_wrappers.hpp" +#include +#include + +#include + +namespace cuvs::neighbors::composite { + +/** + * @brief Merge multiple indices into a single composite index. + * + * This function provides polymorphic merge capability for different index types. + * It delegates to the first index's merge method, which handles the actual merging + * logic based on the index type and merge parameters. + * + * @tparam T Data element type + * @tparam IdxT Index type for vector indices + * @tparam OutputIdxT Output index type + * @param[in] handle RAFT resources for executing operations + * @param[in] params Merge parameters containing strategy and algorithm-specific settings + * @param[in] indices Vector of IndexWrapper pointers to merge + * @return Shared pointer to merged composite index + */ +template +std::shared_ptr> merge( + const raft::resources& handle, + const cuvs::neighbors::merge_params& params, + std::vector>>& indices); + +} // namespace cuvs::neighbors::composite diff --git a/cpp/include/cuvs/neighbors/index_base.hpp b/cpp/include/cuvs/neighbors/index_base.hpp new file mode 100644 index 0000000000..f7cc65c0f0 --- /dev/null +++ b/cpp/include/cuvs/neighbors/index_base.hpp @@ -0,0 +1,103 @@ +/* + * 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. + */ + +#pragma once + +#include +#include +#include + +#include +#include + +namespace cuvs::neighbors { + +/** + * @brief Common interface for ANN index implementations. + * + * Provides a unified abstraction for different approximate nearest neighbor (ANN) index types, + * enabling polymorphic access to core operations such as build, search, size, and metric. + * + * This interface allows: + * - Polymorphic usage of index implementations across the library + * - Easy integration of new index types via a consistent API + * - Support for composite or hierarchical index structures + * - Seamless compatibility between existing and future components + * + * @tparam T Data type of vectors (e.g., float, int8_t) + * @tparam IdxT Index type for dataset rows + * @tparam OutputIdxT Output index type (defaults to IdxT) + */ +template +struct IndexBase { + using value_type = T; + using index_type = IdxT; + using out_index_type = OutputIdxT; + using matrix_index_type = int64_t; + + // Don't allow copying the index for performance reasons + IndexBase() = default; + IndexBase(const IndexBase&) = delete; + IndexBase(IndexBase&&) = default; + auto operator=(const IndexBase&) -> IndexBase& = delete; + auto operator=(IndexBase&&) -> IndexBase& = default; + virtual ~IndexBase() = default; + + /* Future implementation: + virtual void build( + const raft::resources& handle, + const cuvs::neighbors::build_params& params, + raft::device_matrix_view dataset) = 0; + */ + + /** + * @brief Perform approximate nearest neighbor search. + * + * Searches the index for the k-nearest neighbors of each query point. + * The number of neighbors to find is determined by the neighbors matrix extent. + * + * @param[in] handle CUDA resources for executing operations + * @param[in] params Search parameters specific to the index implementation + * @param[in] queries Matrix of query vectors to search for [n_queries, dim] + * @param[out] neighbors Matrix to store neighbor indices [n_queries, k] + * @param[out] distances Matrix to store distances to neighbors [n_queries, k] + * @param[in] filter Optional filter to exclude certain vectors from search results + */ + virtual void search( + const raft::resources& handle, + const cuvs::neighbors::search_params& params, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& filter = + cuvs::neighbors::filtering::none_sample_filter{}) const = 0; + + /** + * @brief Get the number of vectors in the index. + * + * @return Number of indexed vectors + */ + virtual index_type size() const noexcept = 0; + + /** + * @brief Get the distance metric used by the index. + * + * @return Distance metric type (e.g., L2, InnerProduct) + */ + virtual cuvs::distance::DistanceType metric() const noexcept = 0; +}; + +} // namespace cuvs::neighbors diff --git a/cpp/include/cuvs/neighbors/index_wrappers.hpp b/cpp/include/cuvs/neighbors/index_wrappers.hpp new file mode 100644 index 0000000000..fd13f5f1a5 --- /dev/null +++ b/cpp/include/cuvs/neighbors/index_wrappers.hpp @@ -0,0 +1,199 @@ +/* + * 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. + */ + +#pragma once + +#include +#include + +namespace cuvs::neighbors { + +/** + * @brief Intermediate wrapper layer for handling transition between new and old index designs. + * + * This class provides a common interface for wrapping various index implementations + * and handles compatibility issues between different algorithm implementations. + * It serves as an adapter layer that: + * - Provides default implementations for common functionality + * - Handles parameter conversion between different index types + * - Manages post-processing of search results + * - Facilitates gradual migration from old to new index designs + * + * The primary purpose of this wrapper is to help existing index implementations + * transition from their original designs to the new object-oriented polymorphic + * design based on IndexBase. This allows: + * - Legacy index implementations to work with the new polymorphic interface + * - Gradual refactoring of existing code without breaking changes + * - Unified access patterns across different index types + * - Smooth migration path for users adopting the new architecture + * + * By using this wrapper, existing index types (like CAGRA, IVF-PQ, IVF-Flat, HNSW) + * can be adapted to work with the IndexBase interface without requiring immediate + * complete rewrites of their internal implementations. + * + * @tparam T Data element type (e.g., float, int8, uint8) + * @tparam IdxT Index type for vector indices + * @tparam OutputIdxT Output index type, defaults to IdxT + */ +template +class IndexWrapper : public IndexBase { + public: + /** Type definitions inherited from base class */ + using value_type = typename IndexBase::value_type; + using index_type = typename IndexBase::index_type; + using out_index_type = typename IndexBase::out_index_type; + using matrix_index_type = typename IndexBase::matrix_index_type; + + /** Virtual destructor to enable proper cleanup of derived classes */ + virtual ~IndexWrapper() = default; + + /** + * @brief Merge this index with other indices (optional functionality). + * + * This interface provides polymorphic merge capability for index types that support merging. + * The merge strategy and parameters are determined by the specific merge_params implementation. + * Not all index types need to support merging, so this has a default implementation that + * throws an error. + * + * @param[in] handle RAFT resources for executing operations + * @param[in] params Merge parameters containing strategy and algorithm-specific settings + * @param[in] other_indices Vector of other indices to merge with this one + * @return Shared pointer to merged index + */ + virtual std::shared_ptr> merge( + const raft::resources& handle, + const cuvs::neighbors::merge_params& params, + const std::vector>>& + other_indices) const + { + // Default implementation: not supported + RAFT_FAIL("Merge operation not supported for this index type"); + } + + protected: + /** + * @brief Helper method for derived classes to handle parameter conversion. + * + * Derived classes can override this to provide custom parameter handling + * for their specific index types. The default implementation returns + * the parameters unchanged. + * + * @param[in] params Search parameters to convert + * @return Converted search parameters + */ + virtual const search_params& convert_search_params(const search_params& params) const + { + return params; + } +}; + +/** + * @brief Migrating Existing Algorithms to New Polymorphic Index Architecture + * + * To migrate an existing index algorithm (e.g., IVF-PQ, IVF-Flat, HNSW) to the new polymorphic + * IndexBase architecture, follow these steps: + * + * 1. **Create algorithm-specific wrapper header**: + * - Create `cpp/include/cuvs/neighbors/_index_wrapper.hpp` + * - Define `::IndexWrapper` class inheriting from `cuvs::neighbors::IndexWrapper` + * - Place it in the `cuvs::neighbors::` namespace + * - Provide a `make_index_wrapper()` factory function + * + * 2. **Create algorithm-specific wrapper implementation**: + * - Create `cpp/src/neighbors/_index_wrapper.cu` + * - Implement the wrapper methods (search, size, metric, merge if supported) + * - Bridge existing algorithm API to new IndexBase interface + * - Add explicit template instantiations for supported data types + * + * 3. **Include wrapper in main algorithm header**: + * - Add `#include _index_wrapper.hpp>` at the end of + * `cpp/include/cuvs/neighbors/.hpp` + * + * 4. **Update composite merge support** (if merge is supported): + * - Update `cpp/src/neighbors/composite/merge.cpp` to handle the algorithm + * + * Example structure for algorithm "ivf_pq": + * ``` + * cpp/include/cuvs/neighbors/ivf_pq_index_wrapper.hpp: + * namespace cuvs::neighbors::ivf_pq { + * template + * class IndexWrapper : public cuvs::neighbors::IndexWrapper { + * // Bridge existing ivf_pq::index to new interface + * }; + * + * template + * auto make_index_wrapper(ivf_pq::index* index) + * -> std::shared_ptr>; + * } + * + * cpp/src/neighbors/ivf_pq_index_wrapper.cu: + * // Implementation bridging old API to new interface + explicit instantiations + * + * cpp/include/cuvs/neighbors/ivf_pq.hpp: + * // ... existing ivf_pq algorithm code ... + * } // namespace cuvs::neighbors::ivf_pq + * + * #include + * ``` + * + * Usage example (following CAGRA pattern for same-algorithm composite): + * ```cpp + * // 1. Build multiple CAGRA indices on different data partitions + * auto dataset1 = raft::make_device_matrix(res, size1, dim); + * auto dataset2 = raft::make_device_matrix(res, size2, dim); + * auto dataset3 = raft::make_device_matrix(res, size3, dim); + * + * auto cagra_index1 = cuvs::neighbors::cagra::build(res, params, dataset1.view()); + * auto cagra_index2 = cuvs::neighbors::cagra::build(res, params, dataset2.view()); + * auto cagra_index3 = cuvs::neighbors::cagra::build(res, params, dataset3.view()); + * + * // 2. Wrap each index for polymorphic usage + * auto wrapped_index1 = cuvs::neighbors::cagra::make_index_wrapper(&cagra_index1); + * auto wrapped_index2 = cuvs::neighbors::cagra::make_index_wrapper(&cagra_index2); + * auto wrapped_index3 = cuvs::neighbors::cagra::make_index_wrapper(&cagra_index3); + * + * // 3. Merge indices using the composite merge function + * std::vector>> indices; + * indices.push_back(wrapped_index1); + * indices.push_back(wrapped_index2); + * indices.push_back(wrapped_index3); + * + * cuvs::neighbors::cagra::merge_params merge_params; + * auto merged_index = cuvs::neighbors::composite::merge(res, merge_params, indices); + * + * // 4. Search using the merged index + * auto queries = raft::make_device_matrix(res, n_queries, dim); + * auto neighbors = raft::make_device_matrix(res, n_queries, k); + * auto distances = raft::make_device_matrix(res, n_queries, k); + * + * cuvs::neighbors::cagra::search_params search_params; + * merged_index->search(res, search_params, queries.view(), neighbors.view(), distances.view()); + * + * // The merge function automatically: + * // - Merges all 3 CAGRA indices based on the merge strategy + * // - Returns a unified index that can search across all partitions + * // - Handles global addressing and result merging internally + * ``` + * + * Migration benefits: + * - Enables distributed/partitioned indexing with same design + * - Maintains full backward compatibility with existing APIs + * - Allows gradual transition to new architecture + * - Enables composite index functionality for data partitioning + * - Provides unified search interface across multiple index instances + */ + +} // namespace cuvs::neighbors diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index 3919626c80..2e528cec7c 100644 --- a/cpp/src/neighbors/cagra.cuh +++ b/cpp/src/neighbors/cagra.cuh @@ -337,7 +337,6 @@ void search(raft::resources const& res, auto sample_filter_copy = sample_filter; return search_with_filtering( res, params_copy, idx, queries, neighbors, distances, sample_filter_copy); - return; } catch (const std::bad_cast&) { } diff --git a/cpp/src/neighbors/cagra_index_wrapper.cu b/cpp/src/neighbors/cagra_index_wrapper.cu new file mode 100644 index 0000000000..445ed77cec --- /dev/null +++ b/cpp/src/neighbors/cagra_index_wrapper.cu @@ -0,0 +1,110 @@ +/* + * 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 +#include +#include +#include + +namespace cuvs::neighbors::cagra { + +template +IndexWrapper::IndexWrapper(cuvs::neighbors::cagra::index* idx) + : index_(idx) +{ +} + +template +void IndexWrapper::search( + const raft::resources& handle, + const cuvs::neighbors::search_params& params, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& filter) const +{ + auto const& cagra_params = static_cast(params); + cuvs::neighbors::cagra::search( + handle, cagra_params, *index_, queries, neighbors, distances, filter); +} + +template +typename IndexWrapper::index_type IndexWrapper::size() + const noexcept +{ + return index_->size(); +} + +template +cuvs::distance::DistanceType IndexWrapper::metric() const noexcept +{ + return index_->metric(); +} + +template +std::shared_ptr< + cuvs::neighbors::IndexBase::value_type, + typename IndexWrapper::index_type, + typename IndexWrapper::out_index_type>> +IndexWrapper::merge( + const raft::resources& handle, + const cuvs::neighbors::merge_params& params, + const std::vector< + std::shared_ptr>>& + other_indices) const +{ + const auto* cagra_params = dynamic_cast(¶ms); + if (!cagra_params) { RAFT_FAIL("CAGRA IndexWrapper::merge requires cagra::merge_params"); } + + std::vector*> cagra_indices; + cagra_indices.push_back(index_); + + for (const auto& other : other_indices) { + const auto* other_wrapper = dynamic_cast*>(other.get()); + if (!other_wrapper) { + RAFT_FAIL("CAGRA IndexWrapper::merge can only merge with other CAGRA indices"); + } + cagra_indices.push_back(other_wrapper->index_); + } + + if (cagra_params->strategy() == cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_LOGICAL) { + std::vector>> wrappers; + wrappers.reserve(cagra_indices.size()); + for (auto* idx : cagra_indices) { + wrappers.push_back(std::make_shared>(idx)); + } + 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)); + return std::make_shared>(idx); + } + + RAFT_FAIL("Invalid merge strategy"); +} + +template class IndexWrapper; +template class IndexWrapper; +template class IndexWrapper; +template class IndexWrapper; + +template class IndexWrapper; +template class IndexWrapper; +template class IndexWrapper; +template class IndexWrapper; + +} // namespace cuvs::neighbors::cagra diff --git a/cpp/src/neighbors/composite/index.cu b/cpp/src/neighbors/composite/index.cu new file mode 100644 index 0000000000..314ec681f7 --- /dev/null +++ b/cpp/src/neighbors/composite/index.cu @@ -0,0 +1,130 @@ +/* + * Copyright (c) 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. + */ + +#include +#include +#include +#include +#include +#include +#include + +namespace cuvs::neighbors::composite { + +template +void CompositeIndex::search( + const raft::resources& handle, + const cuvs::neighbors::search_params& params, + raft::device_matrix_view queries, + raft::device_matrix_view neighbors, + raft::device_matrix_view distances, + const cuvs::neighbors::filtering::base_filter& filter) const +{ + if (children_.empty()) { + RAFT_FAIL("The composite index is empty!"); + return; + } + + if (children_.size() == 1) { + children_.front()->search(handle, params, queries, neighbors, distances, filter); + return; + } + + size_t num_queries = queries.extent(0); + size_t K = neighbors.extent(1); + size_t num_indices = children_.size(); + size_t buffer_size = num_queries * K * num_indices; + + auto main_stream = raft::resource::get_cuda_stream(handle); + auto tmp_res = raft::resource::get_workspace_resource(handle); + + rmm::device_uvector neighbors_buffer(buffer_size, main_stream, tmp_res); + rmm::device_uvector distances_buffer(buffer_size, main_stream, tmp_res); + + std::vector> temp_neighbors; + std::vector> temp_distances; + + for (size_t i = 0; i < num_indices; i++) { + temp_neighbors.emplace_back(num_queries * K, main_stream, tmp_res); + temp_distances.emplace_back(num_queries * K, main_stream, tmp_res); + } + + raft::resource::wait_stream_pool_on_stream(handle); + + out_index_type offset = 0; + out_index_type stride = K * num_indices; + + for (size_t i = 0; i < num_indices; i++) { + const auto& sub_index = children_[i]; + + auto stream = raft::resource::get_next_usable_stream(handle, i); + + raft::resources stream_pool_handle(handle); + raft::resource::set_cuda_stream(stream_pool_handle, stream); + + auto temp_neighbors_view = + raft::make_device_matrix_view( + temp_neighbors[i].data(), num_queries, K); + auto temp_distances_view = + raft::make_device_matrix_view( + temp_distances[i].data(), num_queries, K); + + sub_index->search( + stream_pool_handle, params, queries, temp_neighbors_view, temp_distances_view, filter); + + if (offset != 0) { + raft::linalg::addScalar(temp_neighbors[i].data(), + temp_neighbors[i].data(), + offset, + temp_neighbors[i].size(), + stream); + } + + raft::copy_matrix( + neighbors_buffer.data() + i * K, stride, temp_neighbors[i].data(), K, K, num_queries, stream); + raft::copy_matrix( + distances_buffer.data() + i * K, stride, temp_distances[i].data(), K, K, num_queries, stream); + + offset += sub_index->size(); + } + raft::resource::sync_stream_pool(handle); + + auto distances_view = raft::make_device_matrix_view( + distances_buffer.data(), num_queries, K * num_indices); + auto neighbors_view = raft::make_device_matrix_view( + neighbors_buffer.data(), num_queries, K * num_indices); + + cuvs::selection::select_k(handle, + distances_view, + neighbors_view, + distances, + neighbors, + cuvs::distance::is_min_close(metric()), + true, // stable_sort + cuvs::selection::SelectAlgo::kAuto); +} + +// Explicit instantiations +template class CompositeIndex; +template class CompositeIndex; +template class CompositeIndex; +template class CompositeIndex; +template class CompositeIndex; +template class CompositeIndex; +template class CompositeIndex; +template class CompositeIndex; + +} // namespace cuvs::neighbors::composite diff --git a/cpp/src/neighbors/composite/merge.cpp b/cpp/src/neighbors/composite/merge.cpp new file mode 100644 index 0000000000..5a4351c548 --- /dev/null +++ b/cpp/src/neighbors/composite/merge.cpp @@ -0,0 +1,83 @@ +/* + * 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 +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cuvs::neighbors::composite { + +/** + * @brief Merge multiple indices into a single composite index. + * + * This function provides polymorphic merge capability for different index types. + * It delegates to the first index's merge method, which handles the actual merging + * logic based on the index type and merge parameters. + * + * @tparam T Data element type + * @tparam IdxT Index type for vector indices + * @tparam OutputIdxT Output index type + * @param[in] handle RAFT resources for executing operations + * @param[in] params Merge parameters containing strategy and algorithm-specific settings + * @param[in] indices Vector of IndexWrapper pointers to merge + * @return Shared pointer to merged composite index + */ +template +std::shared_ptr> merge( + const raft::resources& handle, + const cuvs::neighbors::merge_params& params, + std::vector>>& indices) +{ + if (indices.empty()) { RAFT_FAIL("Cannot merge empty indices vector"); } + + if (indices.size() == 1) { return indices[0]; } + + auto first_index = indices[0]; + std::vector>> other_indices; + for (std::size_t i = 1; i < indices.size(); ++i) { + other_indices.push_back(indices[i]); + } + + // Delegate to the first index's merge method + return first_index->merge(handle, params, other_indices); +} + +#define INSTANTIATE(T, IdxT, OutputIdxT) \ + template std::shared_ptr> \ + merge( \ + const raft::resources&, \ + const cuvs::neighbors::merge_params&, \ + std::vector>>&); + +INSTANTIATE(float, uint32_t, uint32_t); +INSTANTIATE(half, uint32_t, uint32_t); +INSTANTIATE(int8_t, uint32_t, uint32_t); +INSTANTIATE(uint8_t, uint32_t, uint32_t); + +INSTANTIATE(float, uint32_t, int64_t); +INSTANTIATE(half, uint32_t, int64_t); +INSTANTIATE(int8_t, uint32_t, int64_t); +INSTANTIATE(uint8_t, uint32_t, int64_t); + +#undef INSTANTIATE + +} // namespace cuvs::neighbors::composite diff --git a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh index d85866161d..192f7f1632 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 @@ -78,15 +79,13 @@ index merge(raft::resources const& handle, auto merge_dataset = [&](T* dst) { for (cagra_index_t* index : indices) { 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))); + raft::copy_matrix(dst + offset * dim, + dim, + strided_dset->view().data_handle(), + static_cast(stride), + dim, + static_cast(strided_dset->n_rows()), + raft::resource::get_cuda_stream(handle)); offset += IdxT(index->data().n_rows()); } diff --git a/cpp/src/neighbors/detail/cagra/cagra_search.cuh b/cpp/src/neighbors/detail/cagra/cagra_search.cuh index 8ae5445961..79daf2ad6d 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_search.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_search.cuh @@ -26,6 +26,7 @@ #include #include #include +#include #include diff --git a/cpp/tests/neighbors/ann_cagra.cuh b/cpp/tests/neighbors/ann_cagra.cuh index 27812e57c9..4ddc3bd6a7 100644 --- a/cpp/tests/neighbors/ann_cagra.cuh +++ b/cpp/tests/neighbors/ann_cagra.cuh @@ -23,6 +23,8 @@ #include #include +#include +#include #include #include #include @@ -282,6 +284,8 @@ struct AnnCagraInputs { std::optional compression = std::nullopt; std::optional non_owning_memory_buffer_flag = std::nullopt; + cuvs::neighbors::MergeStrategy merge_strategy = + cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL; }; inline ::std::ostream& operator<<(::std::ostream& os, const AnnCagraInputs& p) @@ -298,11 +302,13 @@ inline ::std::ostream& operator<<(::std::ostream& os, const AnnCagraInputs& p) std::vector algo = {"single-cta", "multi_cta", "multi_kernel", "auto"}; std::vector build_algo = {"IVF_PQ", "NN_DESCENT", "ITERATIVE_CAGRA_SEARCH", "AUTO"}; + std::vector merge_strategy = {"PHYSICAL", "LOGICAL"}; os << "{n_queries=" << p.n_queries << ", dataset shape=" << p.n_rows << "x" << p.dim << ", k=" << p.k << ", " << algo.at((int)p.algo) << ", max_queries=" << p.max_queries << ", itopk_size=" << p.itopk_size << ", search_width=" << p.search_width << ", metric=" << metric_str(p.metric) << ", " << (p.host_dataset ? "host" : "device") - << ", build_algo=" << build_algo.at((int)p.build_algo); + << ", build_algo=" << build_algo.at((int)p.build_algo) + << ", merge_logic=" << merge_strategy.at((int)p.merge_strategy); if ((int)p.build_algo == 0 && p.ivf_pq_search_refine_ratio) { os << "(refine_rate=" << *p.ivf_pq_search_refine_ratio << ')'; } @@ -936,6 +942,7 @@ class AnnCagraIndexMergeTest : public ::testing::TestWithParam { } protected: + template void testCagra() { // TODO (tarang-jain): remove when NN Descent index building support InnerProduct. Reference @@ -961,25 +968,25 @@ class AnnCagraIndexMergeTest : public ::testing::TestWithParam { if (ps.n_rows < 8 && ps.build_algo == graph_build_algo::IVF_PQ) 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 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_); + 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); + 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_); @@ -987,7 +994,7 @@ class AnnCagraIndexMergeTest : public ::testing::TestWithParam { { rmm::device_uvector distances_dev(queries_size, stream_); - rmm::device_uvector indices_dev(queries_size, stream_); + rmm::device_uvector indices_dev(queries_size, stream_); { cagra::index_params index_params; @@ -1057,14 +1064,22 @@ class AnnCagraIndexMergeTest : public ::testing::TestWithParam { index0 = cagra::build(handle_, index_params, database0_view); index1 = cagra::build(handle_, index_params, database1_view); }; - std::vector*> indices{&index0, &index1}; + + // Convert traditional CAGRA indices to wrappers for polymorphic usage + std::vector>> + wrapped_indices; + wrapped_indices.push_back( + std::make_shared>(&index0)); + wrapped_indices.push_back( + std::make_shared>(&index1)); + cagra::merge_params merge_params{index_params}; - auto index = cagra::merge(handle_, merge_params, indices); + merge_params.merge_strategy = ps.merge_strategy; 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 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); @@ -1074,8 +1089,11 @@ class AnnCagraIndexMergeTest : public ::testing::TestWithParam { 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); + auto index = cuvs::neighbors::composite::merge( + handle_, merge_params, wrapped_indices); + index->search( + handle_, search_params, 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_); @@ -1136,7 +1154,7 @@ inline std::vector generate_inputs() std::vector inputs = raft::util::itertools::product( {100}, {1000}, - {1, 8, 17}, + {1, 8, 16}, {16}, // k {graph_build_algo::NN_DESCENT, graph_build_algo::ITERATIVE_CAGRA_SEARCH}, // build algo. ITERATIVE_CAGRA_SEARCH is needed to @@ -1149,7 +1167,12 @@ inline std::vector generate_inputs() {cuvs::distance::DistanceType::L2Expanded, cuvs::distance::DistanceType::InnerProduct}, {false}, {true}, - {0.995}); + {0.995}, + {std::optional{std::nullopt}}, + {std::optional{std::nullopt}}, + {std::optional{std::nullopt}}, + {cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL, + cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_LOGICAL}); // Fixed dim, and changing neighbors and query size (output matrix size) auto inputs2 = raft::util::itertools::product( @@ -1172,7 +1195,7 @@ inline std::vector generate_inputs() // Corner cases for small datasets inputs2 = raft::util::itertools::product( {2}, - {3, 5, 31, 32, 64, 101}, + {3, 6, 31, 32, 64, 101}, {1, 10}, {2}, // k {graph_build_algo::IVF_PQ, graph_build_algo::NN_DESCENT}, @@ -1184,7 +1207,12 @@ inline std::vector generate_inputs() {cuvs::distance::DistanceType::L2Expanded}, {false}, {true}, - {0.995}); + {0.995}, + {std::optional{std::nullopt}}, + {std::optional{std::nullopt}}, + {std::optional{std::nullopt}}, + {cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL, + cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_LOGICAL}); inputs.insert(inputs.end(), inputs2.begin(), inputs2.end()); // Varying dim and build algo. @@ -1206,7 +1234,12 @@ inline std::vector generate_inputs() cuvs::distance::DistanceType::BitwiseHamming}, {false}, {true}, - {0.995}); + {0.995}, + {std::optional{std::nullopt}}, + {std::optional{std::nullopt}}, + {std::optional{std::nullopt}}, + {cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL, + cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_LOGICAL}); inputs.insert(inputs.end(), inputs2.begin(), inputs2.end()); // Varying team_size, graph_build_algo @@ -1226,7 +1259,12 @@ inline std::vector generate_inputs() {cuvs::distance::DistanceType::L2Expanded, cuvs::distance::DistanceType::InnerProduct}, {false}, {false}, - {0.995}); + {0.995}, + {std::optional{std::nullopt}}, + {std::optional{std::nullopt}}, + {std::optional{std::nullopt}}, + {cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL, + cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_LOGICAL}); inputs.insert(inputs.end(), inputs2.begin(), inputs2.end()); // Varying graph_build_algo, itopk_size @@ -1246,7 +1284,12 @@ inline std::vector generate_inputs() {cuvs::distance::DistanceType::L2Expanded, cuvs::distance::DistanceType::InnerProduct}, {false}, {true}, - {0.995}); + {0.995}, + {std::optional{std::nullopt}}, + {std::optional{std::nullopt}}, + {std::optional{std::nullopt}}, + {cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL, + cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_LOGICAL}); inputs.insert(inputs.end(), inputs2.begin(), inputs2.end()); // Varying n_rows, host_dataset @@ -1264,7 +1307,12 @@ inline std::vector generate_inputs() {cuvs::distance::DistanceType::L2Expanded, cuvs::distance::DistanceType::InnerProduct}, {false, true}, {false}, - {0.985}); + {0.985}, + {std::optional{std::nullopt}}, + {std::optional{std::nullopt}}, + {std::optional{std::nullopt}}, + {cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL, + cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_LOGICAL}); inputs.insert(inputs.end(), inputs2.begin(), inputs2.end()); // A few PQ configurations. @@ -1283,8 +1331,14 @@ inline std::vector generate_inputs() {cuvs::distance::DistanceType::L2Expanded}, {false}, {true}, - {0.6}); // don't demand high recall without refinement - for (uint32_t pq_len : {2}) { // for now, only pq_len = 2 is supported, more options coming soon + {0.6}, + {std::optional{std::nullopt}}, + {std::optional{std::nullopt}}, + {std::optional{std::nullopt}}, + {cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL, + cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_LOGICAL}); // don't demand high recall + // without refinement + for (uint32_t pq_len : {2}) { // for now, only pq_len = 2 is supported, more options coming soon for (uint32_t vq_n_centers : {100, 1000}) { for (auto input : inputs2) { vpq_params ps{}; @@ -1313,25 +1367,34 @@ inline std::vector generate_inputs() {false, true}, {false}, {0.99}, - {1.0f, 2.0f, 3.0f}); + {1.0f, 2.0f, 3.0f}, + {std::optional{std::nullopt}}, + {std::optional{std::nullopt}}, + {cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL, + cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_LOGICAL}); inputs.insert(inputs.end(), inputs2.begin(), inputs2.end()); // Varying dim, adding non_owning_memory_buffer_flag - inputs2 = - raft::util::itertools::product({100}, - {1000}, - {1, 5, 8, 64, 137, 256, 619, 1024}, // dim - {10}, - {graph_build_algo::IVF_PQ}, - {search_algo::AUTO}, - {10}, - {0}, // team_size - {64}, - {1}, - {cuvs::distance::DistanceType::L2Expanded}, - {false}, - {false}, - {0.995}); + inputs2 = raft::util::itertools::product( + {100}, + {1000}, + {1, 5, 8, 64, 137, 256, 619, 1024}, // dim + {10}, + {graph_build_algo::IVF_PQ}, + {search_algo::AUTO}, + {10}, + {0}, // team_size + {64}, + {1}, + {cuvs::distance::DistanceType::L2Expanded}, + {false}, + {false}, + {0.995}, + {std::optional{std::nullopt}}, + {std::optional{std::nullopt}}, + {std::optional{std::nullopt}}, + {cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL, + cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_LOGICAL}); for (auto input : inputs2) { input.non_owning_memory_buffer_flag = true; inputs.push_back(input); 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 aed372b84c..501550c3b0 100644 --- a/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu +++ b/cpp/tests/neighbors/ann_cagra/test_float_uint32_t.cu @@ -31,7 +31,8 @@ typedef AnnCagraFilterTest AnnCagraFilterTestF_U32; TEST_P(AnnCagraFilterTestF_U32, AnnCagra) { this->testCagra(); } typedef AnnCagraIndexMergeTest AnnCagraIndexMergeTestF_U32; -TEST_P(AnnCagraIndexMergeTestF_U32, AnnCagraIndexMerge) { this->testCagra(); } +TEST_P(AnnCagraIndexMergeTestF_U32, AnnCagraIndexMerge_U32) { this->testCagra(); } +TEST_P(AnnCagraIndexMergeTestF_U32, AnnCagraIndexMerge_I64) { this->testCagra(); } INSTANTIATE_TEST_CASE_P(AnnCagraTest, AnnCagraTestF_U32, ::testing::ValuesIn(inputs)); INSTANTIATE_TEST_CASE_P(AnnCagraAddNodesTest, 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 fdf0ca3bd9..8e44cca410 100644 --- a/cpp/tests/neighbors/ann_cagra/test_half_uint32_t.cu +++ b/cpp/tests/neighbors/ann_cagra/test_half_uint32_t.cu @@ -25,7 +25,8 @@ TEST_P(AnnCagraTestF16_U32, AnnCagra_U32) { this->testCagra(); } TEST_P(AnnCagraTestF16_U32, AnnCagra_I64) { this->testCagra(); } typedef AnnCagraIndexMergeTest AnnCagraIndexMergeTestF16_U32; -TEST_P(AnnCagraIndexMergeTestF16_U32, AnnCagraIndexMerge) { this->testCagra(); } +TEST_P(AnnCagraIndexMergeTestF16_U32, AnnCagraIndexMerge_U32) { this->testCagra(); } +TEST_P(AnnCagraIndexMergeTestF16_U32, AnnCagraIndexMerge_I64) { this->testCagra(); } INSTANTIATE_TEST_CASE_P(AnnCagraTest, AnnCagraTestF16_U32, ::testing::ValuesIn(inputs)); INSTANTIATE_TEST_CASE_P(AnnCagraIndexMergeTest,