Fix filtered CAGRA search returning incorrect recall due to kHostMmap

cpp/bench/ann/src/common/ann_types.hpp

@@ -9,6 +9,7 @@
 
 #include <memory>
 #include <nlohmann/json.hpp>
+#include <optional>
 #include <stdexcept>
 #include <string>
 #include <vector>
@@ -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<MemoryType> filter_memory_type{std::nullopt};
 };
 
 class algo_base {

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<configuration::index>& indices = conf.get_indices();

cpp/bench/ann/src/common/conf.hpp

@@ -44,6 +44,7 @@ class configuration {
     std::string dtype;
 
     std::optional<double> filtering_rate{std::nullopt};
+    std::optional<std::string> 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 =

cpp/bench/ann/src/common/dataset.hpp

@@ -179,7 +179,8 @@ struct dataset {
           std::string query_file,
           std::string distance,
           std::optional<std::string> groundtruth_neighbors_file,
-          std::optional<double> filtering_rate = std::nullopt)
+          std::optional<double> filtering_rate          = std::nullopt,
+          std::optional<std::string> 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<bitset_carrier_type>{filter_bitset_file.value()});
     }
 
     if (groundtruth_neighbors_file.has_value()) {

cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h

@@ -149,12 +149,12 @@ class cuvs_cagra : public algo<T>, 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<T, IdxT>::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<T, int64_t>(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<size_t>(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<size_t>(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<T, int64_t>(handle_, 0, 0);
+      auto dataset_view = raft::make_device_strided_matrix_view<const T, int64_t>(
+        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<T, int64_t>(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<const T, int64_t>(
-      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<const T, int64_t>(
+        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<T, IdxT>::load(const std::string& file)
 template <typename T, typename IdxT>
 std::unique_ptr<algo<T>> cuvs_cagra<T, IdxT>::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<std::shared_ptr<cuvs::neighbors::cagra::index<T, IdxT>>>;
+    new (&sub_indices_) SubVec{};
+  }
   return std::make_unique<cuvs_cagra<T, IdxT>>(std::cref(*this));  // use copy constructor
 }