diff --git a/cpp/bench/ann/src/common/ann_types.hpp b/cpp/bench/ann/src/common/ann_types.hpp index bd669dff78..2562ef58a3 100644 --- a/cpp/bench/ann/src/common/ann_types.hpp +++ b/cpp/bench/ann/src/common/ann_types.hpp @@ -9,6 +9,7 @@ #include #include +#include #include #include #include @@ -78,6 +79,8 @@ struct algo_property { MemoryType dataset_memory_type; // neighbors/distances should have same memory type as queries MemoryType query_memory_type; + // filter bitset memory type; defaults to dataset_memory_type if not explicitly set + std::optional filter_memory_type{std::nullopt}; }; class algo_base { diff --git a/cpp/bench/ann/src/common/benchmark.hpp b/cpp/bench/ann/src/common/benchmark.hpp index a588b1e2a6..4ef6cc076d 100644 --- a/cpp/bench/ann/src/common/benchmark.hpp +++ b/cpp/bench/ann/src/common/benchmark.hpp @@ -102,6 +102,9 @@ inline auto parse_algo_property(algo_property prop, const nlohmann::json& conf) if (conf.contains("query_memory_type")) { prop.query_memory_type = parse_memory_type(conf.at("query_memory_type")); } + if (conf.contains("filter_memory_type")) { + prop.filter_memory_type = parse_memory_type(conf.at("filter_memory_type")); + } return prop; }; @@ -263,7 +266,8 @@ void bench_search(::benchmark::State& state, } try { a->set_search_param(*search_param, - dataset->filter_bitset(current_algo_props->dataset_memory_type)); + dataset->filter_bitset(current_algo_props->filter_memory_type.value_or( + current_algo_props->dataset_memory_type))); } catch (const std::exception& ex) { state.SkipWithError("An error occurred setting search parameters: " + std::string(ex.what())); return; @@ -545,7 +549,8 @@ void dispatch_benchmark(std::string cmdline, query_file, dataset_conf.distance, gt_file, - search_mode ? dataset_conf.filtering_rate : std::nullopt); + search_mode ? dataset_conf.filtering_rate : std::nullopt, + dataset_conf.filter_bitset_file); ::benchmark::AddCustomContext("dataset", dataset_conf.name); ::benchmark::AddCustomContext("distance", dataset_conf.distance); std::vector& indices = conf.get_indices(); diff --git a/cpp/bench/ann/src/common/conf.hpp b/cpp/bench/ann/src/common/conf.hpp index afc7bc0a1f..fa213c106c 100644 --- a/cpp/bench/ann/src/common/conf.hpp +++ b/cpp/bench/ann/src/common/conf.hpp @@ -44,6 +44,7 @@ class configuration { std::string dtype; std::optional filtering_rate{std::nullopt}; + std::optional filter_bitset_file{std::nullopt}; }; [[nodiscard]] inline auto get_dataset_conf() const -> const dataset_conf& @@ -89,6 +90,9 @@ class configuration { if (conf.contains("filtering_rate")) { dataset_conf_.filtering_rate.emplace(conf.at("filtering_rate")); } + if (conf.contains("filter_bitset_file")) { + dataset_conf_.filter_bitset_file = combine_path(data_prefix, conf.at("filter_bitset_file")); + } if (conf.contains("groundtruth_neighbors_file")) { dataset_conf_.groundtruth_neighbors_file = diff --git a/cpp/bench/ann/src/common/dataset.hpp b/cpp/bench/ann/src/common/dataset.hpp index 4dc43c343c..04954dfaf6 100644 --- a/cpp/bench/ann/src/common/dataset.hpp +++ b/cpp/bench/ann/src/common/dataset.hpp @@ -179,7 +179,8 @@ struct dataset { std::string query_file, std::string distance, std::optional groundtruth_neighbors_file, - std::optional filtering_rate = std::nullopt) + std::optional filtering_rate = std::nullopt, + std::optional filter_bitset_file = std::nullopt) : name_{std::move(name)}, distance_{std::move(distance)}, base_set_{base_file, subset_first_row, subset_size}, @@ -200,6 +201,8 @@ struct dataset { bitset_size, 1.0 - filtering_rate.value()); filter_bitset_.emplace(std::move(bitset_blob)); + } else if (filter_bitset_file.has_value()) { + filter_bitset_.emplace(blob{filter_bitset_file.value()}); } if (groundtruth_neighbors_file.has_value()) { diff --git a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h index 87111e4761..e336472725 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h @@ -149,12 +149,12 @@ class cuvs_cagra : public algo, public algo_gpu { return !search_params_.persistent; } - // to enable dataset access from GPU memory [[nodiscard]] auto get_preference() const -> algo_property override { algo_property property; property.dataset_memory_type = MemoryType::kHostMmap; property.query_memory_type = MemoryType::kDevice; + property.filter_memory_type = MemoryType::kDevice; return property; } void save(const std::string& file) const override; @@ -313,20 +313,38 @@ void cuvs_cagra::set_search_param(const search_param_base& param, if (sp.dataset_mem != dataset_mem_ || need_dataset_update_) { dataset_mem_ = sp.dataset_mem; - // First free up existing memory - *dataset_ = raft::make_device_matrix(handle_, 0, 0); - index_->update_dataset(handle_, make_const_mdspan(dataset_->view())); + // When the benchmark framework provides the dataset on device (kDevice) and no padding is + // needed for alignment, skip the redundant copy_with_padding() allocation. For a 100M-scale + // dataset, copy_with_padding() would double the device memory requirement (e.g. 38.4 GB x2), + // causing OOM. Instead, use the existing device allocation directly. + size_t padded_dim = raft::round_up_safe(this->dim_ * sizeof(T), 16) / sizeof(T); + bool data_on_device = raft::get_device_for_address(input_dataset_v_->data_handle()) >= 0 && + sp.dataset_mem == AllocatorType::kDevice; + bool padding_needed = padded_dim != static_cast(this->dim_); + + if (data_on_device && !padding_needed) { + // Data is already in device memory and no padding is needed — update the index directly. + RAFT_LOG_DEBUG("dataset already on device, skipping copy_with_padding"); + *dataset_ = raft::make_device_matrix(handle_, 0, 0); + auto dataset_view = raft::make_device_strided_matrix_view( + input_dataset_v_->data_handle(), input_dataset_v_->extent(0), this->dim_, this->dim_); + index_->update_dataset(handle_, dataset_view); + } else { + // First free up existing memory + *dataset_ = raft::make_device_matrix(handle_, 0, 0); + index_->update_dataset(handle_, make_const_mdspan(dataset_->view())); - // Allocate space using the correct memory resource. - RAFT_LOG_DEBUG("moving dataset to new memory space: %s", - allocator_to_string(dataset_mem_).c_str()); + // Allocate space using the correct memory resource. + RAFT_LOG_DEBUG("moving dataset to new memory space: %s", + allocator_to_string(dataset_mem_).c_str()); - auto mr = get_mr(dataset_mem_); - cuvs::neighbors::cagra::detail::copy_with_padding(handle_, *dataset_, *input_dataset_v_, mr); + auto mr = get_mr(dataset_mem_); + cuvs::neighbors::cagra::detail::copy_with_padding(handle_, *dataset_, *input_dataset_v_, mr); - auto dataset_view = raft::make_device_strided_matrix_view( - dataset_->data_handle(), dataset_->extent(0), this->dim_, dataset_->extent(1)); - index_->update_dataset(handle_, dataset_view); + auto dataset_view = raft::make_device_strided_matrix_view( + dataset_->data_handle(), dataset_->extent(0), this->dim_, dataset_->extent(1)); + index_->update_dataset(handle_, dataset_view); + } need_dataset_update_ = false; needs_dynamic_batcher_update = true; @@ -452,6 +470,13 @@ void cuvs_cagra::load(const std::string& file) template std::unique_ptr> cuvs_cagra::copy() { + // sub_indices_ can get corrupted when dataset_memory_type=kDevice triggers copy_with_padding + // during set_search_param. For single-index CAGRA, sub_indices_ is always logically empty, + // so we reset it via placement new before copying to avoid bad_array_new_length. + if (index_params_.num_dataset_splits <= 1) { + using SubVec = std::vector>>; + new (&sub_indices_) SubVec{}; + } return std::make_unique>(std::cref(*this)); // use copy constructor }