Skip to content

Commit

Permalink
Fix 0 recall issue in raft_cagra_hnswlib ANN benchmark (#2369)
Browse files Browse the repository at this point in the history
`raft_cagra` wrapper stopped including the dataset in the index to save memory, but this adversely affected `raft_cagra_hnswlib` wrapper because the dataset needed to be included in the index. The need for inclusion of the dataset is because we need the dataset to be serialized when writing to the `hnswlib` format.

Authors:
  - Divye Gala (https://github.com/divyegala)

Approvers:
  - Corey J. Nolet (https://github.com/cjnolet)

URL: #2369
  • Loading branch information
divyegala authored Jun 27, 2024
1 parent ad4e543 commit 36f77a1
Show file tree
Hide file tree
Showing 2 changed files with 10 additions and 4 deletions.
2 changes: 1 addition & 1 deletion cpp/bench/ann/src/raft/raft_cagra_hnswlib_wrapper.h
Original file line number Diff line number Diff line change
Expand Up @@ -31,7 +31,7 @@ class RaftCagraHnswlib : public ANN<T>, public AnnGPU {

RaftCagraHnswlib(Metric metric, int dim, const BuildParam& param, int concurrent_searches = 1)
: ANN<T>(metric, dim),
cagra_build_{metric, dim, param, concurrent_searches},
cagra_build_{metric, dim, param, concurrent_searches, true},
// HnswLib param values don't matter since we don't build with HnswLib
hnswlib_search_{metric, dim, typename HnswLib<T>::BuildParam{50, 100}}
{
Expand Down
12 changes: 9 additions & 3 deletions cpp/bench/ann/src/raft/raft_cagra_wrapper.h
Original file line number Diff line number Diff line change
Expand Up @@ -72,11 +72,16 @@ class RaftCagra : public ANN<T>, public AnnGPU {
std::optional<raft::neighbors::ivf_pq::search_params> ivf_pq_search_params = std::nullopt;
};

RaftCagra(Metric metric, int dim, const BuildParam& param, int concurrent_searches = 1)
RaftCagra(Metric metric,
int dim,
const BuildParam& param,
int concurrent_searches = 1,
bool shall_include_dataset = false)
: ANN<T>(metric, dim),
index_params_(param),
dimension_(dim),
need_dataset_update_(true),
shall_include_dataset_(shall_include_dataset),
dataset_(std::make_shared<raft::device_matrix<T, int64_t, row_major>>(
std::move(make_device_matrix<T, int64_t>(handle_, 0, 0)))),
graph_(std::make_shared<raft::device_matrix<IdxT, int64_t, row_major>>(
Expand Down Expand Up @@ -135,6 +140,7 @@ class RaftCagra : public ANN<T>, public AnnGPU {
float refine_ratio_;
BuildParam index_params_;
bool need_dataset_update_;
bool shall_include_dataset_;
raft::neighbors::cagra::search_params search_params_;
std::shared_ptr<raft::neighbors::cagra::index<T, IdxT>> index_;
int dimension_;
Expand All @@ -161,7 +167,7 @@ void RaftCagra<T, IdxT>::build(const T* dataset, size_t nrow)
auto& params = index_params_.cagra_params;

// Do include the compressed dataset for the CAGRA-Q
bool shall_include_dataset = params.compression.has_value();
bool include_dataset = params.compression.has_value() || shall_include_dataset_;

index_ = std::make_shared<raft::neighbors::cagra::index<T, IdxT>>(
std::move(raft::neighbors::cagra::detail::build(handle_,
Expand All @@ -171,7 +177,7 @@ void RaftCagra<T, IdxT>::build(const T* dataset, size_t nrow)
index_params_.ivf_pq_refine_rate,
index_params_.ivf_pq_build_params,
index_params_.ivf_pq_search_params,
shall_include_dataset)));
include_dataset)));
}

inline std::string allocator_to_string(AllocatorType mem_type)
Expand Down

0 comments on commit 36f77a1

Please sign in to comment.