diff --git a/conda/environments/bench_ann_cuda-118_arch-x86_64.yaml b/conda/environments/bench_ann_cuda-118_arch-x86_64.yaml index b7344c822..56ca20ad6 100644 --- a/conda/environments/bench_ann_cuda-118_arch-x86_64.yaml +++ b/conda/environments/bench_ann_cuda-118_arch-x86_64.yaml @@ -26,6 +26,8 @@ dependencies: - gcc_linux-64=11.* - glog>=0.6.0 - h5py>=3.8.0 +- libaio +- libboost-devel - libcublas-dev=11.11.3.6 - libcublas=11.11.3.6 - libcurand-dev=10.3.0.86 @@ -37,6 +39,7 @@ dependencies: - libcuvs==25.2.*,>=0.0.0a0 - librmm==25.2.*,>=0.0.0a0 - matplotlib +- mkl-devel - nccl>=2.19 - ninja - nlohmann_json>=3.11.2 diff --git a/conda/environments/bench_ann_cuda-125_arch-x86_64.yaml b/conda/environments/bench_ann_cuda-125_arch-x86_64.yaml index 5d1dd8fc7..6ef31d3e8 100644 --- a/conda/environments/bench_ann_cuda-125_arch-x86_64.yaml +++ b/conda/environments/bench_ann_cuda-125_arch-x86_64.yaml @@ -27,6 +27,8 @@ dependencies: - gcc_linux-64=11.* - glog>=0.6.0 - h5py>=3.8.0 +- libaio +- libboost-devel - libcublas-dev - libcurand-dev - libcusolver-dev @@ -34,6 +36,7 @@ dependencies: - libcuvs==25.2.*,>=0.0.0a0 - librmm==25.2.*,>=0.0.0a0 - matplotlib +- mkl-devel - nccl>=2.19 - ninja - nlohmann_json>=3.11.2 diff --git a/conda/recipes/cuvs-bench-cpu/meta.yaml b/conda/recipes/cuvs-bench-cpu/meta.yaml index 016df56be..282277a87 100644 --- a/conda/recipes/cuvs-bench-cpu/meta.yaml +++ b/conda/recipes/cuvs-bench-cpu/meta.yaml @@ -47,6 +47,9 @@ requirements: - benchmark - fmt {{ fmt_version }} - glog {{ glog_version }} + - libaio + - libboost-devel + - mkl-devel # [linux64] - nlohmann_json {{ nlohmann_json_version }} - openblas - python diff --git a/conda/recipes/cuvs-bench/meta.yaml b/conda/recipes/cuvs-bench/meta.yaml index 0681a1038..2065c7315 100644 --- a/conda/recipes/cuvs-bench/meta.yaml +++ b/conda/recipes/cuvs-bench/meta.yaml @@ -72,7 +72,10 @@ requirements: - libcublas-dev {% endif %} - glog {{ glog_version }} + - libaio + - libboost-devel - libcuvs {{ version }} + - mkl-devel # [linux64] - nlohmann_json {{ nlohmann_json_version }} - openblas # rmm is needed to determine if package is gpu-enabled diff --git a/cpp/bench/ann/CMakeLists.txt b/cpp/bench/ann/CMakeLists.txt index 144cd3048..81ba0b98e 100644 --- a/cpp/bench/ann/CMakeLists.txt +++ b/cpp/bench/ann/CMakeLists.txt @@ -32,6 +32,12 @@ option(CUVS_ANN_BENCH_USE_CUVS_BRUTE_FORCE "Include cuVS brute force knn in benc option(CUVS_ANN_BENCH_USE_CUVS_CAGRA_HNSWLIB "Include cuVS CAGRA with HNSW search in benchmark" ON) option(CUVS_ANN_BENCH_USE_HNSWLIB "Include hnsw algorithm in benchmark" ON) option(CUVS_ANN_BENCH_USE_GGNN "Include ggnn algorithm in benchmark" OFF) +option(CUVS_ANN_BENCH_USE_DISKANN "Include DISKANN search in benchmark" ON) +option(CUVS_ANN_BENCH_USE_CUVS_VAMANA "Include cuVS Vamana with DiskANN search in benchmark" ON) +if(CMAKE_SYSTEM_PROCESSOR MATCHES "(ARM|arm|aarch64)") + set(CUVS_ANN_BENCH_USE_DISKANN OFF) + set(CUVS_ANN_BENCH_USE_CUVS_VAMANA OFF) +endif() option(CUVS_ANN_BENCH_USE_CUVS_MG "Include cuVS ann mg algorithm in benchmark" ${BUILD_MG_ALGOS}) option(CUVS_ANN_BENCH_SINGLE_EXE "Make a single executable with benchmark as shared library modules" OFF @@ -57,6 +63,7 @@ if(BUILD_CPU_ONLY) set(CUVS_ANN_BENCH_USE_GGNN OFF) set(CUVS_KNN_BENCH_USE_CUVS_BRUTE_FORCE OFF) set(CUVS_ANN_BENCH_USE_CUVS_MG OFF) + set(CUVS_ANN_BENCH_USE_CUVS_VAMANA OFF) else() set(CUVS_FAISS_ENABLE_GPU ON) endif() @@ -69,6 +76,7 @@ if(CUVS_ANN_BENCH_USE_CUVS_IVF_PQ OR CUVS_ANN_BENCH_USE_CUVS_CAGRA_HNSWLIB OR CUVS_KNN_BENCH_USE_CUVS_BRUTE_FORCE OR CUVS_ANN_BENCH_USE_CUVS_MG + OR CUVS_ANN_BENCH_USE_CUVS_VAMANA ) set(CUVS_ANN_BENCH_USE_CUVS ON) endif() @@ -90,6 +98,10 @@ if(CUVS_ANN_BENCH_USE_FAISS) include(cmake/thirdparty/get_faiss) endif() +if(CUVS_ANN_BENCH_USE_DISKANN OR CUVS_ANN_BENCH_USE_CUVS_VAMANA) + include(cmake/thirdparty/get_diskann) +endif() + # ################################################################################################## # * Target function ------------------------------------------------------------- @@ -294,6 +306,19 @@ if(CUVS_ANN_BENCH_USE_GGNN) ) endif() +if(CUVS_ANN_BENCH_USE_DISKANN) + ConfigureAnnBench( + NAME DISKANN_MEMORY PATH src/diskann/diskann_benchmark.cpp LINKS diskann::diskann aio + ) + ConfigureAnnBench( + NAME DISKANN_SSD PATH src/diskann/diskann_benchmark.cpp LINKS diskann::diskann aio + ) +endif() + +if(CUVS_ANN_BENCH_USE_CUVS_VAMANA) + ConfigureAnnBench(NAME CUVS_VAMANA PATH src/cuvs/cuvs_vamana.cu LINKS cuvs diskann::diskann aio) +endif() + # ################################################################################################## # * Dynamically-loading ANN_BENCH executable ------------------------------------------------------- if(CUVS_ANN_BENCH_SINGLE_EXE) diff --git a/cpp/bench/ann/src/common/benchmark.hpp b/cpp/bench/ann/src/common/benchmark.hpp index 06e1e27af..7f507cd22 100644 --- a/cpp/bench/ann/src/common/benchmark.hpp +++ b/cpp/bench/ann/src/common/benchmark.hpp @@ -135,6 +135,12 @@ void bench_build(::benchmark::State& state, } } + if (index.algo == "diskann_ssd") { + make_sure_parent_dir_exists(index.file); + index.build_param["dataset_file"] = dataset->base_filename(); + index.build_param["path_to_index"] = index.file; + } + std::unique_ptr> algo; try { algo = create_algo(index.algo, dataset->distance(), dataset->dim(), index.build_param); @@ -144,7 +150,8 @@ void bench_build(::benchmark::State& state, const auto algo_property = parse_algo_property(algo->get_preference(), index.build_param); - const T* base_set = dataset->base_set(algo_property.dataset_memory_type); + const T* base_set = nullptr; + if (index.algo != "diskann_ssd") base_set = dataset->base_set(algo_property.dataset_memory_type); std::size_t index_size = dataset->base_set_size(); cuda_timer gpu_timer{algo}; @@ -223,7 +230,12 @@ void bench_search(::benchmark::State& state, const T* query_set = nullptr; - if (!file_exists(index.file)) { + std::string filename; + if (index.algo != "diskann_ssd") + filename = index.file; + else + filename = index.file + "_disk.index"; + if (!file_exists(filename)) { state.SkipWithError("Index file is missing. Run the benchmark in the build mode first."); return; } diff --git a/cpp/bench/ann/src/common/dataset.hpp b/cpp/bench/ann/src/common/dataset.hpp index 49020fe36..c3f565f61 100644 --- a/cpp/bench/ann/src/common/dataset.hpp +++ b/cpp/bench/ann/src/common/dataset.hpp @@ -114,6 +114,8 @@ class bin_file { } } + std::string file() const { return file_; } + private: void check_suffix(); void open_file() const; @@ -253,10 +255,11 @@ class dataset { auto name() const -> std::string { return name_; } auto distance() const -> std::string { return distance_; } - virtual auto dim() const -> int = 0; - virtual auto max_k() const -> uint32_t = 0; - virtual auto base_set_size() const -> size_t = 0; - virtual auto query_set_size() const -> size_t = 0; + virtual auto dim() const -> int = 0; + virtual auto max_k() const -> uint32_t = 0; + virtual auto base_set_size() const -> size_t = 0; + virtual auto query_set_size() const -> size_t = 0; + virtual auto base_filename() const -> std::string = 0; // load data lazily, so don't pay the overhead of reading unneeded set // e.g. don't load base set when searching @@ -424,6 +427,7 @@ class bin_dataset : public dataset { auto max_k() const -> uint32_t override; auto base_set_size() const -> size_t override; auto query_set_size() const -> size_t override; + std::string base_filename() const override; private: void load_base_set() const; @@ -541,4 +545,10 @@ void bin_dataset::map_base_set() const this->mapped_base_set_ = base_file_.map(); } +template +std::string bin_dataset::base_filename() const +{ + return base_file_.file(); +} + } // namespace cuvs::bench diff --git a/cpp/bench/ann/src/cuvs/cuvs_vamana.cu b/cpp/bench/ann/src/cuvs/cuvs_vamana.cu new file mode 100644 index 000000000..5b9dde859 --- /dev/null +++ b/cpp/bench/ann/src/cuvs/cuvs_vamana.cu @@ -0,0 +1,93 @@ +/* + * 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 "../common/ann_types.hpp" +#include "cuvs_vamana_wrapper.h" + +#include +#include +#include + +namespace cuvs::bench { + +template +void parse_build_param(const nlohmann::json& conf, + typename cuvs::bench::cuvs_vamana::build_param& param) +{ + if (conf.contains("graph_degree")) { param.graph_degree = conf.at("graph_degree"); } + if (conf.contains("visited_size")) { param.visited_size = conf.at("visited_size"); } + if (conf.contains("alpha")) { param.alpha = conf.at("alpha"); } +} + +template +void parse_search_param(const nlohmann::json& conf, + typename cuvs::bench::cuvs_vamana::search_param& param) +{ + if (conf.contains("L_search")) { param.L_search = conf.at("L_search"); } + if (conf.contains("num_threads")) { param.num_threads = conf.at("num_threads"); } +} + +template +auto create_algo(const std::string& algo_name, + const std::string& distance, + int dim, + const nlohmann::json& conf) -> std::unique_ptr> +{ + [[maybe_unused]] cuvs::bench::Metric metric = parse_metric(distance); + std::unique_ptr> a; + + if constexpr (std::is_same_v or std::is_same_v) { + if (algo_name == "cuvs_vamana") { + typename cuvs::bench::cuvs_vamana::build_param param; + parse_build_param(conf, param); + a = std::make_unique>(metric, dim, param); + } + } + + if (!a) { throw std::runtime_error("invalid algo: '" + algo_name + "'"); } + + return a; +} + +template +auto create_search_param(const std::string& algo_name, const nlohmann::json& conf) + -> std::unique_ptr::search_param> +{ + if (algo_name == "cuvs_vamana") { + auto param = std::make_unique::search_param>(); + parse_search_param(conf, *param); + return param; + } + + throw std::runtime_error("invalid algo: '" + algo_name + "'"); +} + +} // namespace cuvs::bench + +REGISTER_ALGO_INSTANCE(float); + +#ifdef ANN_BENCH_BUILD_MAIN +#include "../common/benchmark.hpp" +/* +[NOTE] Dear developer, + +Please don't modify the content of the `main` function; this will make the behavior of the benchmark +executable differ depending on the cmake flags and will complicate the debugging. In particular, +don't try to setup an RMM memory resource here; it will anyway be modified by the memory resource +set on per-algorithm basis. For example, see `cuvs/cuvs_ann_bench_utils.h`. +*/ +int main(int argc, char** argv) { return cuvs::bench::run_main(argc, argv); } +#endif diff --git a/cpp/bench/ann/src/cuvs/cuvs_vamana_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_vamana_wrapper.h new file mode 100644 index 000000000..76536e83d --- /dev/null +++ b/cpp/bench/ann/src/cuvs/cuvs_vamana_wrapper.h @@ -0,0 +1,124 @@ +/* + * 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. + */ +#pragma once + +#include "../common/ann_types.hpp" +#include "../diskann/diskann_wrapper.h" +#include "cuvs_ann_bench_utils.h" +#include +#include + +#include +#include +#include + +namespace cuvs::bench { + +template +class cuvs_vamana : public algo, public algo_gpu { + public: + using build_param = cuvs::neighbors::experimental::vamana::index_params; + using search_param_base = typename algo::search_param; + using search_param = typename diskann_memory::search_param; + + cuvs_vamana(Metric metric, int dim, const build_param& param); + + void build(const T* dataset, size_t nrow) final; + + void set_search_param(const search_param_base& param) override; + + void search(const T* queries, + int batch_size, + int k, + algo_base::index_type* neighbors, + float* distances) const override; + + [[nodiscard]] auto get_sync_stream() const noexcept -> cudaStream_t override + { + return handle_.get_sync_stream(); + } + + // to enable dataset access from GPU memory + [[nodiscard]] auto get_preference() const -> algo_property override + { + algo_property property; + property.dataset_memory_type = MemoryType::kDevice; + property.query_memory_type = MemoryType::kHost; + return property; + } + + void save(const std::string& file) const override; + void load(const std::string&) override; + std::unique_ptr> copy() override { return std::make_unique>(*this); } + + private: + std::shared_ptr> vamana_index_; + std::shared_ptr> diskann_memory_search_; + configured_raft_resources handle_{}; + build_param vamana_index_params_; +}; + +template +cuvs_vamana::cuvs_vamana(Metric metric, int dim, const build_param& param) + : algo(metric, dim) +{ + this->vamana_index_params_ = param; + diskann_memory_search_ = std::make_shared>( + metric, dim, typename diskann_memory::build_param{param.graph_degree, param.visited_size}); +} + +template +void cuvs_vamana::build(const T* dataset, size_t nrow) +{ + auto dataset_view_host = raft::make_mdspan( + dataset, raft::make_extents(nrow, this->dim_)); + auto dataset_view_device = raft::make_mdspan( + dataset, raft::make_extents(nrow, this->dim_)); + bool dataset_is_on_host = raft::get_device_for_address(dataset) == -1; + + vamana_index_ = std::make_shared>( + std::move(dataset_is_on_host ? cuvs::neighbors::experimental::vamana::build( + handle_, vamana_index_params_, dataset_view_host) + : cuvs::neighbors::experimental::vamana::build( + handle_, vamana_index_params_, dataset_view_device))); +} + +template +void cuvs_vamana::set_search_param(const search_param_base& param_) +{ + diskann_memory_search_->set_search_param(param_); +} + +template +void cuvs_vamana::save(const std::string& file) const +{ + cuvs::neighbors::experimental::vamana::serialize(handle_, file, *vamana_index_); +} + +template +void cuvs_vamana::load(const std::string& file) +{ + diskann_memory_search_->load(file); +} + +template +void cuvs_vamana::search( + const T* queries, int batch_size, int k, algo_base::index_type* neighbors, float* distances) const +{ + diskann_memory_search_->search(queries, batch_size, k, neighbors, distances); +} + +} // namespace cuvs::bench diff --git a/cpp/bench/ann/src/diskann/diskann_benchmark.cpp b/cpp/bench/ann/src/diskann/diskann_benchmark.cpp new file mode 100644 index 000000000..ca8a94048 --- /dev/null +++ b/cpp/bench/ann/src/diskann/diskann_benchmark.cpp @@ -0,0 +1,131 @@ +/* + * 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 "../common/ann_types.hpp" +#include "diskann_wrapper.h" + +#define JSON_DIAGNOSTICS 1 +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace cuvs::bench { + +template +void parse_build_param(const nlohmann::json& conf, + typename cuvs::bench::diskann_memory::build_param& param) +{ + param.R = conf.at("R"); + if (conf.contains("L_build")) { param.L_build = conf.at("L_build"); } + if (conf.contains("alpha")) { param.num_threads = conf.at("alpha"); } + if (conf.contains("num_threads")) { param.num_threads = conf.at("num_threads"); } +} + +template +void parse_build_param(const nlohmann::json& conf, + typename cuvs::bench::diskann_ssd::build_param& param) +{ + param.R = conf.at("R"); + if (conf.contains("L_build")) { param.L_build = conf.at("L_build"); } + if (conf.contains("alpha")) { param.num_threads = conf.at("alpha"); } + if (conf.contains("num_threads")) { param.num_threads = conf.at("num_threads"); } + if (conf.contains("QD")) { param.QD = conf.at("QD"); } + if (conf.contains("dataset_file")) { param.dataset_file = conf.at("dataset_file"); } + if (conf.contains("path_to_index")) { param.path_to_index = conf.at("path_to_index"); } +} + +template +void parse_search_param(const nlohmann::json& conf, + typename cuvs::bench::diskann_memory::search_param& param) +{ + param.L_search = conf.at("L_search"); + param.num_threads = conf.at("num_threads"); +} + +template +void parse_search_param(const nlohmann::json& conf, + typename cuvs::bench::diskann_ssd::search_param& param) +{ + param.L_search = conf.at("L_search"); + param.num_threads = conf.at("num_threads"); + if (conf.contains("num_nodes_to_cache")) { + param.num_nodes_to_cache = conf.at("num_nodes_to_cache"); + } + if (conf.contains("beam_width")) { param.beam_width = conf.at("beam_width"); } +} + +template class Algo> +std::unique_ptr> make_algo(cuvs::bench::Metric metric, + int dim, + const nlohmann::json& conf) +{ + typename Algo::build_param param; + parse_build_param(conf, param); + return std::make_unique>(metric, dim, param); +} + +template +auto create_algo(const std::string& algo_name, + const std::string& distance, + int dim, + const nlohmann::json& conf) -> std::unique_ptr> +{ + cuvs::bench::Metric metric = parse_metric(distance); + std::unique_ptr> a; + + if constexpr (std::is_same_v || std::is_same_v || + std::is_same_v) { + if (algo_name == "diskann_memory") { + a = make_algo(metric, dim, conf); + } else if (algo_name == "diskann_ssd") { + a = make_algo(metric, dim, conf); + } + } + if (!a) { throw std::runtime_error("invalid algo: '" + algo_name + "'"); } + + return a; +} + +template +std::unique_ptr::search_param> create_search_param( + const std::string& algo_name, const nlohmann::json& conf) +{ + if (algo_name == "diskann_memory") { + auto param = std::make_unique::search_param>(); + parse_search_param(conf, *param); + return param; + } else if (algo_name == "diskann_ssd") { + auto param = std::make_unique::search_param>(); + parse_search_param(conf, *param); + return param; + } + throw std::runtime_error("invalid algo: '" + algo_name + "'"); +} + +}; // namespace cuvs::bench + +REGISTER_ALGO_INSTANCE(float); + +#ifdef ANN_BENCH_BUILD_MAIN +#include "../common/benchmark.hpp" +int main(int argc, char** argv) { return cuvs::bench::run_main(argc, argv); } +#endif \ No newline at end of file diff --git a/cpp/bench/ann/src/diskann/diskann_wrapper.h b/cpp/bench/ann/src/diskann/diskann_wrapper.h new file mode 100644 index 000000000..79f207c81 --- /dev/null +++ b/cpp/bench/ann/src/diskann/diskann_wrapper.h @@ -0,0 +1,335 @@ +/* + * 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. + */ +#pragma once + +#include "../common/ann_types.hpp" + +#include + +#include +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace cuvs::bench { + +diskann::Metric parse_metric_to_diskann(cuvs::bench::Metric metric) +{ + if (metric == cuvs::bench::Metric::kInnerProduct) { + return diskann::Metric::INNER_PRODUCT; + } else if (metric == cuvs::bench::Metric::kEuclidean) { + return diskann::Metric::L2; + } else { + throw std::runtime_error("currently only inner product and L2 supported for benchmarking"); + } +} + +template +class diskann_memory : public algo { + public: + struct build_param { + uint32_t R; + uint32_t L_build; + uint32_t build_pq_bytes = 0; + float alpha = 1.2; + int num_threads = omp_get_num_procs(); + }; + + using search_param_base = typename algo::search_param; + struct search_param : public search_param_base { + uint32_t L_search; + uint32_t num_threads = omp_get_num_procs(); + // Mode metric_objective; + }; + + diskann_memory(Metric metric, int dim, const build_param& param); + + void build(const T* dataset, size_t nrow) override; + + void set_search_param(const search_param_base& param) override; + + void search(const T* queries, + int batch_size, + int k, + algo_base::index_type* indices, + float* distances) const override; + + void save(const std::string& path_to_index) const override; + void load(const std::string& path_to_index) override; + diskann_memory(const diskann_memory& other) = default; + std::unique_ptr> copy() override { return std::make_unique>(*this); } + + [[nodiscard]] auto get_preference() const -> algo_property override + { + algo_property property; + property.dataset_memory_type = MemoryType::kHost; + property.query_memory_type = MemoryType::kHost; + return property; + } + + private: + std::shared_ptr diskann_index_write_params_{nullptr}; + uint32_t max_points_; + uint32_t build_pq_bytes_ = 0; + int num_threads_; + uint32_t L_search_; + Mode bench_mode_; + int num_search_threads_; + std::string index_path_prefix_; + std::shared_ptr> mem_index_{nullptr}; + void initialize_index_(size_t max_points); +}; + +template +diskann_memory::diskann_memory(Metric metric, int dim, const build_param& param) + : algo(metric, dim) +{ + assert(this->dim_ > 0); + num_threads_ = param.num_threads; + diskann_index_write_params_ = std::make_shared( + diskann::IndexWriteParametersBuilder(param.L_build, param.R) + .with_filter_list_size(0) + .with_alpha(param.alpha) + .with_saturate_graph(false) + .with_num_threads(param.num_threads) + .build()); +} + +template +void diskann_memory::initialize_index_(size_t max_points) +{ + this->mem_index_ = std::make_shared>(parse_metric_to_diskann(this->metric_), + this->dim_, + max_points, + diskann_index_write_params_, + nullptr, + 0, + false, + false, + false, + build_pq_bytes_ > 0, + build_pq_bytes_, + false, + false); +} +template +void diskann_memory::build(const T* dataset, size_t nrow) +{ + initialize_index_(nrow); + mem_index_->build(dataset, nrow, std::vector()); +} + +template +void diskann_memory::set_search_param(const search_param_base& param_) +{ + auto param = dynamic_cast(param_); + L_search_ = param.L_search; + num_search_threads_ = param.num_threads; + + // only latency mode supported. Use the num_threads search param to run search with multiple + // threads + bench_mode_ = Mode::kLatency; + + // Create a pool if multiple query threads have been set and the pool hasn't been created already + initialize_index_(0); + this->mem_index_->load(index_path_prefix_.c_str(), num_search_threads_, L_search_); +} + +template +void diskann_memory::search( + const T* queries, int batch_size, int k, algo_base::index_type* indices, float* distances) const +{ +#pragma omp parallel for schedule(dynamic, 1) + for (int i = 0; i < batch_size; i++) { + mem_index_->search(queries + i * this->dim_, + static_cast(k), + L_search_, + reinterpret_cast(indices + i * k), + distances + i * k); + } +} + +template +void diskann_memory::save(const std::string& path_to_index) const +{ + this->mem_index_->save(path_to_index.c_str()); +} + +template +void diskann_memory::load(const std::string& path_to_index) +{ + // only save the index path prefix here + index_path_prefix_ = path_to_index; +} + +template +class diskann_ssd : public algo { + public: + struct build_param { + uint32_t R; + uint32_t L_build; + uint32_t build_pq_bytes = 0; + float alpha = 1.2; + int num_threads = omp_get_num_procs(); + uint32_t QD = 192; + std::string dataset_file = ""; + std::string path_to_index = ""; + }; + using search_param_base = typename algo::search_param; + + struct search_param : public search_param_base { + uint32_t L_search; + uint32_t num_threads = omp_get_num_procs() / 2; + uint32_t num_nodes_to_cache = 10000; + int beam_width = 2; + // Mode metric_objective; + }; + + diskann_ssd(Metric metric, int dim, const build_param& param); + + void build(const T* dataset, size_t nrow) override; + + void set_search_param(const search_param_base& param) override; + + void search(const T* queries, + int batch_size, + int k, + algo_base::index_type* neighbors, + float* distances) const override; + + void save(const std::string& path_to_index) const override; + void load(const std::string& path_to_index) override; + diskann_ssd(const diskann_ssd& other) = default; + std::unique_ptr> copy() override { return std::make_unique>(*this); } + + [[nodiscard]] auto get_preference() const -> algo_property override + { + algo_property property; + property.dataset_memory_type = MemoryType::kHost; + property.query_memory_type = MemoryType::kHost; + return property; + } + + private: + std::string index_build_params_str; + std::shared_ptr> p_flash_index_; + int beam_width_; + uint32_t num_nodes_to_cache_; + + // in-memory index params + uint32_t build_pq_bytes_ = 0; + uint32_t max_points_; + // for safe scratch space allocs, set the default to half the number of procs for loading the + // index. User must ensure that the number of search threads is less than or equal to this value + int num_search_threads_ = omp_get_num_procs() / 2; + // L_search is hardcoded to the maximum visited list size in the search params. This default is + // for loading the index + uint32_t L_search_ = 384; + Mode bench_mode_; + std::string base_file_; + std::string index_path_prefix_; + std::shared_ptr reader = nullptr; +}; + +template +diskann_ssd::diskann_ssd(Metric metric, int dim, const build_param& param) : algo(metric, dim) +{ + // Currently set the indexing RAM budget and the search RAM budget to max value to avoid sharding + uint32_t build_dram_budget = std::numeric_limits::max(); + uint32_t search_dram_budget = std::numeric_limits::max(); + index_build_params_str = + std::string(std::to_string(param.R)) + " " + std::string(std::to_string(param.L_build)) + " " + + std::string(std::to_string(search_dram_budget)) + " " + + std::string(std::to_string(build_dram_budget)) + " " + + std::string(std::to_string(param.num_threads)) + " " + std::string(std::to_string(false)) + + " " + std::string(std::to_string(false)) + " " + std::string(std::to_string(0)) + " " + + std::string(std::to_string(param.QD)); + base_file_ = param.dataset_file; + index_path_prefix_ = param.path_to_index; +} + +template +void diskann_ssd::build(const T* dataset, size_t nrow) +{ + diskann::build_disk_index(base_file_.c_str(), + index_path_prefix_.c_str(), + index_build_params_str.c_str(), + parse_metric_to_diskann(this->metric_), + false, + std::string(""), + false, + std::string(""), + std::string(""), + 0, + 0); +} + +template +void diskann_ssd::set_search_param(const search_param_base& param_) +{ + auto param = dynamic_cast(param_); + L_search_ = param.L_search; + num_search_threads_ = param.num_threads; + num_nodes_to_cache_ = param.num_nodes_to_cache; + beam_width_ = param.beam_width; + + // only latency mode supported with thread pool + bench_mode_ = Mode::kLatency; +} + +template +void diskann_ssd::search( + const T* queries, int batch_size, int k, algo_base::index_type* neighbors, float* distances) const +{ +#pragma omp parallel for schedule(dynamic, 1) + for (int64_t i = 0; i < (int64_t)batch_size; i++) { + p_flash_index_->cached_beam_search(queries + (i * this->dim_), + static_cast(k), + L_search_, + reinterpret_cast(neighbors + i * k), + distances + i * k, + beam_width_, + false, + nullptr); + } +} + +template +void diskann_ssd::save(const std::string& path_to_index) const +{ + // Nothing to do here. Index already saved in build stage. +} + +template +void diskann_ssd::load(const std::string& path_to_index) +{ + reader.reset(new LinuxAlignedFileReader()); + p_flash_index_ = + std::make_shared>(reader, parse_metric_to_diskann(this->metric_)); + int result = p_flash_index_->load(num_search_threads_, path_to_index.c_str()); + std::vector node_list; + p_flash_index_->cache_bfs_levels(num_nodes_to_cache_, node_list); + p_flash_index_->load_cache_list(node_list); + node_list.clear(); + node_list.shrink_to_fit(); +} +}; // namespace cuvs::bench diff --git a/cpp/cmake/patches/diskann.diff b/cpp/cmake/patches/diskann.diff new file mode 100644 index 000000000..466b24d22 --- /dev/null +++ b/cpp/cmake/patches/diskann.diff @@ -0,0 +1,227 @@ +diff --git a/CMakeLists.txt b/CMakeLists.txt +index 3d3d2b8..3079d12 100644 +--- a/CMakeLists.txt ++++ b/CMakeLists.txt +@@ -145,62 +145,14 @@ if (MSVC) + "${DISKANN_MKL_LIB_PATH}/mkl_intel_thread.lib") + else() + # expected path for manual intel mkl installs +- set(POSSIBLE_OMP_PATHS "/opt/intel/oneapi/compiler/latest/linux/compiler/lib/intel64_lin/libiomp5.so;/usr/lib/x86_64-linux-gnu/libiomp5.so;/opt/intel/lib/intel64_lin/libiomp5.so") +- foreach(POSSIBLE_OMP_PATH ${POSSIBLE_OMP_PATHS}) +- if (EXISTS ${POSSIBLE_OMP_PATH}) +- get_filename_component(OMP_PATH ${POSSIBLE_OMP_PATH} DIRECTORY) +- endif() +- endforeach() +- +- if(NOT OMP_PATH) +- message(FATAL_ERROR "Could not find Intel OMP in standard locations; use -DOMP_PATH to specify the install location for your environment") +- endif() +- link_directories(${OMP_PATH}) +- +- set(POSSIBLE_MKL_LIB_PATHS "/opt/intel/oneapi/mkl/latest/lib/intel64/libmkl_core.so;/usr/lib/x86_64-linux-gnu/libmkl_core.so;/opt/intel/mkl/lib/intel64/libmkl_core.so") +- foreach(POSSIBLE_MKL_LIB_PATH ${POSSIBLE_MKL_LIB_PATHS}) +- if (EXISTS ${POSSIBLE_MKL_LIB_PATH}) +- get_filename_component(MKL_PATH ${POSSIBLE_MKL_LIB_PATH} DIRECTORY) +- endif() +- endforeach() +- +- set(POSSIBLE_MKL_INCLUDE_PATHS "/opt/intel/oneapi/mkl/latest/include;/usr/include/mkl;/opt/intel/mkl/include/;") +- foreach(POSSIBLE_MKL_INCLUDE_PATH ${POSSIBLE_MKL_INCLUDE_PATHS}) +- if (EXISTS ${POSSIBLE_MKL_INCLUDE_PATH}) +- set(MKL_INCLUDE_PATH ${POSSIBLE_MKL_INCLUDE_PATH}) +- endif() +- endforeach() +- if(NOT MKL_PATH) +- message(FATAL_ERROR "Could not find Intel MKL in standard locations; use -DMKL_PATH to specify the install location for your environment") +- elseif(NOT MKL_INCLUDE_PATH) +- message(FATAL_ERROR "Could not find Intel MKL in standard locations; use -DMKL_INCLUDE_PATH to specify the install location for headers for your environment") +- endif() +- if (EXISTS ${MKL_PATH}/libmkl_def.so.2) +- set(MKL_DEF_SO ${MKL_PATH}/libmkl_def.so.2) +- elseif(EXISTS ${MKL_PATH}/libmkl_def.so) +- set(MKL_DEF_SO ${MKL_PATH}/libmkl_def.so) +- else() +- message(FATAL_ERROR "Despite finding MKL, libmkl_def.so was not found in expected locations.") +- endif() +- link_directories(${MKL_PATH}) +- include_directories(${MKL_INCLUDE_PATH}) ++ find_package(MKL CONFIG REQUIRED) ++ include_directories($) ++ link_libraries($) + + # compile flags and link libraries + add_compile_options(-m64 -Wl,--no-as-needed) + if (NOT PYBIND) + link_libraries(mkl_intel_ilp64 mkl_intel_thread mkl_core iomp5 pthread m dl) +- else() +- # static linking for python so as to minimize customer dependency issues +- link_libraries( +- ${MKL_PATH}/libmkl_intel_ilp64.a +- ${MKL_PATH}/libmkl_intel_thread.a +- ${MKL_PATH}/libmkl_core.a +- ${MKL_DEF_SO} +- iomp5 +- pthread +- m +- dl +- ) + endif() + endif() + +@@ -286,7 +238,7 @@ if(MSVC) + set(CMAKE_ARCHIVE_OUTPUT_DIRECTORY_RELEASE ${PROJECT_SOURCE_DIR}/x64/Release) + else() + set(ENV{TCMALLOC_LARGE_ALLOC_REPORT_THRESHOLD} 500000000000) +- set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -mavx2 -mfma -msse2 -ftree-vectorize -fno-builtin-malloc -fno-builtin-calloc -fno-builtin-realloc -fno-builtin-free -fopenmp -fopenmp-simd -funroll-loops -Wfatal-errors -DUSE_AVX2") ++ set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -mavx2 -mfma -msse2 -ftree-vectorize -fno-builtin-malloc -fno-builtin-calloc -fno-builtin-realloc -fno-builtin-free -fopenmp -fopenmp-simd -funroll-loops -Wfatal-errors -DUSE_AVX2 -fno-finite-math-only -laio") + set(CMAKE_CXX_FLAGS_DEBUG "${CMAKE_CXX_FLAGS_DEBUG} -g -DDEBUG") + if (NOT PYBIND) + set(CMAKE_CXX_FLAGS_RELEASE "${CMAKE_CXX_FLAGS_RELEASE} -DNDEBUG -Ofast") +@@ -300,10 +252,6 @@ else() + endif() + + add_subdirectory(src) +-if (NOT PYBIND) +- add_subdirectory(apps) +- add_subdirectory(apps/utils) +-endif() + + if (UNIT_TEST) + enable_testing() +diff --git a/include/distance.h b/include/distance.h +index f3b1de2..d4da72e 100644 +--- a/include/distance.h ++++ b/include/distance.h +@@ -77,6 +77,7 @@ class DistanceCosineInt8 : public Distance + DistanceCosineInt8() : Distance(diskann::Metric::COSINE) + { + } ++ using Distance::compare; + DISKANN_DLLEXPORT virtual float compare(const int8_t *a, const int8_t *b, uint32_t length) const; + }; + +@@ -86,6 +87,7 @@ class DistanceL2Int8 : public Distance + DistanceL2Int8() : Distance(diskann::Metric::L2) + { + } ++ using Distance::compare; + DISKANN_DLLEXPORT virtual float compare(const int8_t *a, const int8_t *b, uint32_t size) const; + }; + +@@ -96,6 +98,7 @@ class AVXDistanceL2Int8 : public Distance + AVXDistanceL2Int8() : Distance(diskann::Metric::L2) + { + } ++ using Distance::compare; + DISKANN_DLLEXPORT virtual float compare(const int8_t *a, const int8_t *b, uint32_t length) const; + }; + +@@ -105,6 +108,7 @@ class DistanceCosineFloat : public Distance + DistanceCosineFloat() : Distance(diskann::Metric::COSINE) + { + } ++ using Distance::compare; + DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t length) const; + }; + +@@ -114,7 +118,7 @@ class DistanceL2Float : public Distance + DistanceL2Float() : Distance(diskann::Metric::L2) + { + } +- ++ using Distance::compare; + #ifdef _WINDOWS + DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t size) const; + #else +@@ -128,6 +132,7 @@ class AVXDistanceL2Float : public Distance + AVXDistanceL2Float() : Distance(diskann::Metric::L2) + { + } ++ using Distance::compare; + DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t length) const; + }; + +@@ -146,6 +151,7 @@ class SlowDistanceCosineUInt8 : public Distance + SlowDistanceCosineUInt8() : Distance(diskann::Metric::COSINE) + { + } ++ using Distance::compare; + DISKANN_DLLEXPORT virtual float compare(const uint8_t *a, const uint8_t *b, uint32_t length) const; + }; + +@@ -155,6 +161,7 @@ class DistanceL2UInt8 : public Distance + DistanceL2UInt8() : Distance(diskann::Metric::L2) + { + } ++ using Distance::compare; + DISKANN_DLLEXPORT virtual float compare(const uint8_t *a, const uint8_t *b, uint32_t size) const; + }; + +@@ -198,6 +205,7 @@ class AVXDistanceInnerProductFloat : public Distance + AVXDistanceInnerProductFloat() : Distance(diskann::Metric::INNER_PRODUCT) + { + } ++ using Distance::compare; + DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t length) const; + }; + +@@ -213,6 +221,7 @@ class AVXNormalizedCosineDistanceFloat : public Distance + AVXNormalizedCosineDistanceFloat() : Distance(diskann::Metric::COSINE) + { + } ++ using Distance::compare; + DISKANN_DLLEXPORT virtual float compare(const float *a, const float *b, uint32_t length) const + { + // Inner product returns negative values to indicate distance. +diff --git a/include/utils.h b/include/utils.h +index d3af5c3..417af31 100644 +--- a/include/utils.h ++++ b/include/utils.h +@@ -29,6 +29,7 @@ typedef int FileHandle; + #include "types.h" + #include "tag_uint128.h" + #include ++#include + + #ifdef EXEC_ENV_OLS + #include "content_buf.h" +diff --git a/src/index.cpp b/src/index.cpp +index bf93344..9d8336c 100644 +--- a/src/index.cpp ++++ b/src/index.cpp +@@ -17,9 +17,7 @@ + #include "gperftools/malloc_extension.h" + #endif + +-#ifdef _WINDOWS + #include +-#endif + + #include "index.h" + +diff --git a/src/partition.cpp b/src/partition.cpp +index 570d45c..fb54cbf 100644 +--- a/src/partition.cpp ++++ b/src/partition.cpp +@@ -21,9 +21,7 @@ + #include "parameters.h" + #include "memory_mapper.h" + #include "partition.h" +-#ifdef _WINDOWS + #include +-#endif + + // block size for reading/ processing large files and matrices in blocks + #define BLOCK_SIZE 5000000 +diff --git a/src/pq_flash_index.cpp b/src/pq_flash_index.cpp +index d9ad506..145a978 100644 +--- a/src/pq_flash_index.cpp ++++ b/src/pq_flash_index.cpp +@@ -8,6 +8,7 @@ + #include "pq_scratch.h" + #include "pq_flash_index.h" + #include "cosine_similarity.h" ++#include + + #ifdef _WINDOWS + #include "windows_aligned_file_reader.h" diff --git a/cpp/cmake/patches/diskann_override.json b/cpp/cmake/patches/diskann_override.json new file mode 100644 index 000000000..c83898548 --- /dev/null +++ b/cpp/cmake/patches/diskann_override.json @@ -0,0 +1,16 @@ +{ + "packages" : { + "diskann" : { + "version": "0.7.0", + "git_url": "https://github.com/microsoft/DiskANN.git", + "git_tag": "main", + "patches" : [ + { + "file" : "${current_json_dir}/diskann.diff", + "issue" : "Correct compilation issues", + "fixed_in" : "" + } + ] + } + } +} diff --git a/cpp/cmake/thirdparty/get_diskann.cmake b/cpp/cmake/thirdparty/get_diskann.cmake new file mode 100644 index 000000000..8cea20d5c --- /dev/null +++ b/cpp/cmake/thirdparty/get_diskann.cmake @@ -0,0 +1,49 @@ +#============================================================================= +# 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. +#============================================================================= + +function(find_and_configure_diskann) + set(oneValueArgs VERSION REPOSITORY PINNED_TAG) + cmake_parse_arguments(PKG "${options}" "${oneValueArgs}" + "${multiValueArgs}" ${ARGN} ) + + include(${rapids-cmake-dir}/cpm/package_override.cmake) + set(patch_dir "${CMAKE_CURRENT_FUNCTION_LIST_DIR}/../patches") + rapids_cpm_package_override("${patch_dir}/diskann_override.json") + + include("${rapids-cmake-dir}/cpm/detail/package_details.cmake") + rapids_cpm_package_details(diskann version repository tag shallow exclude) + + include("${rapids-cmake-dir}/cpm/detail/generate_patch_command.cmake") + rapids_cpm_generate_patch_command(diskann ${version} patch_command) + + rapids_cpm_find(diskann ${version} + GLOBAL_TARGETS diskann + CPM_ARGS + OPTIONS + "PYBIND OFF" + "UNIT_TEST OFF" + "RESTAPI OFF" + "PORTABLE OFF") + + include("${rapids-cmake-dir}/cpm/detail/display_patch_status.cmake") + rapids_cpm_display_patch_status(diskann) + + if(NOT TARGET diskann::diskann) + target_include_directories(diskann INTERFACE "$") + add_library(diskann::diskann ALIAS diskann) + endif() +endfunction() +find_and_configure_diskann() \ No newline at end of file diff --git a/cpp/include/cuvs/neighbors/vamana.hpp b/cpp/include/cuvs/neighbors/vamana.hpp index bec17937f..634f6fdb1 100644 --- a/cpp/include/cuvs/neighbors/vamana.hpp +++ b/cpp/include/cuvs/neighbors/vamana.hpp @@ -258,15 +258,18 @@ auto build(raft::resources const& handle, void serialize(raft::resources const& handle, const std::string& file_prefix, - const cuvs::neighbors::experimental::vamana::index& index); + const cuvs::neighbors::experimental::vamana::index& index, + bool include_dataset = true); void serialize(raft::resources const& handle, const std::string& file_prefix, - const cuvs::neighbors::experimental::vamana::index& index); + const cuvs::neighbors::experimental::vamana::index& index, + bool include_dataset = true); void serialize(raft::resources const& handle, const std::string& file_prefix, - const cuvs::neighbors::experimental::vamana::index& index); + const cuvs::neighbors::experimental::vamana::index& index, + bool include_dataset = true); /** * @} diff --git a/cpp/src/neighbors/detail/vamana/vamana_serialize.cuh b/cpp/src/neighbors/detail/vamana/vamana_serialize.cuh index a554464f6..ad9d9d770 100644 --- a/cpp/src/neighbors/detail/vamana/vamana_serialize.cuh +++ b/cpp/src/neighbors/detail/vamana/vamana_serialize.cuh @@ -44,13 +44,15 @@ namespace cuvs::neighbors::experimental::vamana::detail { * @param[in] res the raft resource handle * @param[in] file_name the path and name of the DiskAN index file generated * @param[in] index_ VAMANA index + * @param[in] include_dataset whether to include the dataset in the serialized output * */ template void serialize(raft::resources const& res, const std::string& file_name, - const index& index_) + const index& index_, + bool include_dataset) { // Write graph to first index file (format from MSFT DiskANN OSS) std::ofstream index_of(file_name, std::ios::out | std::ios::binary); @@ -75,6 +77,7 @@ void serialize(raft::resources const& res, d_graph.data_handle(), d_graph.size(), raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); size_t total_edges = 0; size_t num_sparse = 0; @@ -115,6 +118,42 @@ void serialize(raft::resources const& res, index_of.close(); if (!index_of) { RAFT_FAIL("Error writing output %s", file_name.c_str()); } + + // try allocating a buffer for the dataset on host + try { + const cuvs::neighbors::strided_dataset* strided_dataset = + dynamic_cast*>( + const_cast*>(&index_.data())); + if (strided_dataset == nullptr) { + RAFT_LOG_DEBUG("dynamic_cast to strided_dataset failed"); + } else { + auto h_dataset = + raft::make_host_matrix(strided_dataset->n_rows(), strided_dataset->dim()); + raft::copy(h_dataset.data_handle(), + strided_dataset->view().data_handle(), + strided_dataset->n_rows() * strided_dataset->dim(), + raft::resource::get_cuda_stream(res)); + std::string dataset_file = file_name + ".data"; + std::ofstream dataset_of(dataset_file, std::ios::out | std::ios::binary); + if (!dataset_of) { RAFT_FAIL("Cannot open file %s", dataset_file.c_str()); } + size_t dataset_file_offset = 0; + int size = static_cast(index_.size()); + int dim = static_cast(index_.dim()); + dataset_of.seekp(dataset_file_offset, dataset_of.beg); + dataset_of.write((char*)&size, sizeof(int)); + dataset_of.write((char*)&dim, sizeof(int)); + for (int i = 0; i < size; i++) { + dataset_of.write((char*)(h_dataset.data_handle() + i * h_dataset.extent(1)), + dim * sizeof(T)); + } + dataset_of.close(); + if (!dataset_of) { RAFT_FAIL("Error writing output %s", dataset_file.c_str()); } + } + } catch (std::bad_alloc& e) { + RAFT_LOG_INFO("Failed to serialize dataset"); + } catch (raft::logic_error& e) { + RAFT_LOG_INFO("Failed to serialize dataset"); + } } } // namespace cuvs::neighbors::experimental::vamana::detail diff --git a/cpp/src/neighbors/vamana.cuh b/cpp/src/neighbors/vamana.cuh index 9b9e8d271..f06b0e427 100644 --- a/cpp/src/neighbors/vamana.cuh +++ b/cpp/src/neighbors/vamana.cuh @@ -94,7 +94,7 @@ void serialize(raft::resources const& res, const std::string& file_prefix, const index& index_) { - cuvs::neighbors::experimental::vamana::detail::build(res, file_prefix, index_); + cuvs::neighbors::experimental::vamana::detail::serialize(res, file_prefix, index_); } /** @} */ // end group vamana diff --git a/cpp/src/neighbors/vamana_serialize.cuh b/cpp/src/neighbors/vamana_serialize.cuh index a49d267b3..f2cd06dec 100644 --- a/cpp/src/neighbors/vamana_serialize.cuh +++ b/cpp/src/neighbors/vamana_serialize.cuh @@ -28,10 +28,11 @@ namespace cuvs::neighbors::experimental::vamana { #define CUVS_INST_VAMANA_SERIALIZE(DTYPE) \ void serialize(raft::resources const& handle, \ const std::string& file_prefix, \ - const cuvs::neighbors::experimental::vamana::index& index_) \ + const cuvs::neighbors::experimental::vamana::index& index_, \ + bool include_dataset) \ { \ cuvs::neighbors::experimental::vamana::detail::serialize( \ - handle, file_prefix, index_); \ + handle, file_prefix, index_, include_dataset); \ }; /** @} */ // end group vamana diff --git a/cpp/test/CMakeLists.txt b/cpp/test/CMakeLists.txt index 4d13daaed..aa1a4cda0 100644 --- a/cpp/test/CMakeLists.txt +++ b/cpp/test/CMakeLists.txt @@ -236,7 +236,7 @@ if(BUILD_TESTS) NAME SPARSE_TEST PATH sparse/cluster/cluster_solvers.cu sparse/cluster/eigen_solvers.cu sparse/cluster/spectral.cu GPUS 1 PERCENT 100 ) - + ConfigureTest( NAME PREPROCESSING_TEST PATH preprocessing/scalar_quantization.cu GPUS 1 PERCENT 100 ) diff --git a/dependencies.yaml b/dependencies.yaml index eca97d2f5..638b94be0 100644 --- a/dependencies.yaml +++ b/dependencies.yaml @@ -477,6 +477,18 @@ dependencies: - benchmark>=1.8.2 - openblas - libcuvs==25.2.*,>=0.0.0a0 + specific: + - output_types: conda + matrices: + - matrix: + arch: x86_64 + packages: + - mkl-devel + - libaio + - libboost-devel + - matrix: + arch: aarch64 + packages: null bench_python: common: - output_types: [conda, pyproject, requirements] diff --git a/python/cuvs_bench/cuvs_bench/config/algorithms.yaml b/python/cuvs_bench/cuvs_bench/config/algorithms.yaml index 357517933..a5f79ff78 100644 --- a/python/cuvs_bench/cuvs_bench/config/algorithms.yaml +++ b/python/cuvs_bench/cuvs_bench/config/algorithms.yaml @@ -49,3 +49,12 @@ hnswlib: cuvs_cagra_hnswlib: executable: CUVS_CAGRA_HNSWLIB_ANN_BENCH requires_gpu: true +diskann_memory: + executable: DISKANN_MEMORY_ANN_BENCH + requires_gpu: false +diskann_ssd: + executable: DISKANN_SSD_ANN_BENCH + requires_gpu: false +cuvs_vamana: + executable: CUVS_VAMANA_ANN_BENCH + requires_gpu: true diff --git a/python/cuvs_bench/cuvs_bench/config/algos/constraints/__init__.py b/python/cuvs_bench/cuvs_bench/config/algos/constraints/__init__.py index de05bd752..dd521bec2 100644 --- a/python/cuvs_bench/cuvs_bench/config/algos/constraints/__init__.py +++ b/python/cuvs_bench/cuvs_bench/config/algos/constraints/__init__.py @@ -99,3 +99,24 @@ def faiss_gpu_ivf_pq_search(params, build_params, k, batch_size): def hnswlib_search(params, build_params, k, batch_size): if "ef" in params: return params["ef"] >= k + + +############################################################################### +# DiskANN constraints # +############################################################################### + + +def diskann_memory_build(params, dim): + ret = True + if "R" in params and "L_build" in params: + ret = params["R"] <= params["L_build"] + return ret + + +def diskann_ssd_build(params, dim): + ret = True + if "R" in params and "L_build" in params: + ret = params["R"] <= params["L_build"] + if "QD" in params: + ret = params["QD"] <= dim + return ret diff --git a/python/cuvs_bench/cuvs_bench/config/algos/cuvs_vamana.yaml b/python/cuvs_bench/cuvs_bench/config/algos/cuvs_vamana.yaml new file mode 100644 index 000000000..5252f818c --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/config/algos/cuvs_vamana.yaml @@ -0,0 +1,10 @@ +name: cuvs_vamana +groups: + base: + build: + graph_degree: [32, 64] + visited_size: [128, 256] + alpha: [1.2] + search: + L_search: [10, 20, 30, 40, 50, 100, 200, 300] + num_threads: [32] diff --git a/python/cuvs_bench/cuvs_bench/config/algos/diskann_memory.yaml b/python/cuvs_bench/cuvs_bench/config/algos/diskann_memory.yaml new file mode 100644 index 000000000..faf122465 --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/config/algos/diskann_memory.yaml @@ -0,0 +1,13 @@ +name: diskann_memory +constraints: + build: cuvs_bench.config.algos.constraints.diskann_memory_build +groups: + base: + build: + R: [64, 96] + L_build: [128, 256, 384] + alpha: [1.2] + num_threads: [32] + search: + num_threads: [32] + L_search: [10, 20, 30, 40, 50, 100, 200, 300] diff --git a/python/cuvs_bench/cuvs_bench/config/algos/diskann_ssd.yaml b/python/cuvs_bench/cuvs_bench/config/algos/diskann_ssd.yaml new file mode 100644 index 000000000..40afc6fce --- /dev/null +++ b/python/cuvs_bench/cuvs_bench/config/algos/diskann_ssd.yaml @@ -0,0 +1,13 @@ +name: diskann_ssd +constraints: + build: cuvs_bench.config.algos.constraints.diskann_ssd_build +groups: + base: + build: + R: [64, 96] + L_build: [128, 256, 384] + QD: [192] + num_threads: [32] + search: + L_search: [10, 20, 30, 40, 50, 100, 200, 300] + num_threads: [32]