From 53e7982c2127df1439012af773f08a57ea2a93c9 Mon Sep 17 00:00:00 2001 From: Divye Gala Date: Thu, 27 Jun 2024 11:20:19 -0400 Subject: [PATCH 1/3] Fix `ef` setting in HNSW wrapper (#2367) Closes #2363 Bugs fixed: 1. Setting `ef` in search, it was not being set at all before 2. `from_cagra` used a hard coded filename to serialize CAGRA graph and deserialize to HNSW graph. The PR changes the hardcoded filename to a random string so that multiple graphs may be converted concurrently cc @Presburger thank you for reporting these bugs Authors: - Divye Gala (https://github.com/divyegala) Approvers: - Corey J. Nolet (https://github.com/cjnolet) URL: https://github.com/rapidsai/raft/pull/2367 --- cpp/include/raft/neighbors/detail/hnsw.hpp | 1 + cpp/include/raft/neighbors/detail/hnsw_types.hpp | 5 +++++ cpp/include/raft/neighbors/hnsw.hpp | 2 +- cpp/include/raft/neighbors/hnsw_types.hpp | 5 +++++ cpp/src/raft_runtime/neighbors/hnsw.cpp | 8 +++++++- python/pylibraft/pylibraft/neighbors/hnsw.pyx | 6 ++++-- 6 files changed, 23 insertions(+), 4 deletions(-) diff --git a/cpp/include/raft/neighbors/detail/hnsw.hpp b/cpp/include/raft/neighbors/detail/hnsw.hpp index 5deee3c2ba..bd4e6608de 100644 --- a/cpp/include/raft/neighbors/detail/hnsw.hpp +++ b/cpp/include/raft/neighbors/detail/hnsw.hpp @@ -53,6 +53,7 @@ void search(raft::resources const& res, raft::host_matrix_view neighbors, raft::host_matrix_view distances) { + idx.set_ef(params.ef); auto const* hnswlib_index = reinterpret_cast::type> const*>( idx.get_index()); diff --git a/cpp/include/raft/neighbors/detail/hnsw_types.hpp b/cpp/include/raft/neighbors/detail/hnsw_types.hpp index 9d35effd1a..8d601f59ae 100644 --- a/cpp/include/raft/neighbors/detail/hnsw_types.hpp +++ b/cpp/include/raft/neighbors/detail/hnsw_types.hpp @@ -93,6 +93,11 @@ struct index_impl : index { */ auto get_index() const -> void const* override { return appr_alg_.get(); } + /** + @brief Set ef for search + */ + void set_ef(int ef) const override { appr_alg_->ef_ = ef; } + private: std::unique_ptr::type>> appr_alg_; std::unique_ptr::type>> space_; diff --git a/cpp/include/raft/neighbors/hnsw.hpp b/cpp/include/raft/neighbors/hnsw.hpp index 964c3ffacd..ee3f61e550 100644 --- a/cpp/include/raft/neighbors/hnsw.hpp +++ b/cpp/include/raft/neighbors/hnsw.hpp @@ -35,7 +35,7 @@ namespace raft::neighbors::hnsw { /** * @brief Construct an hnswlib base-layer-only index from a CAGRA index - * NOTE: 1. This method uses the filesystem to write the CAGRA index in `/tmp/cagra_index.bin` + * NOTE: 1. This method uses the filesystem to write the CAGRA index in `/tmp/.bin` * before reading it as an hnswlib index, then deleting the temporary file. * 2. This function is only offered as a compiled symbol in `libraft.so` * diff --git a/cpp/include/raft/neighbors/hnsw_types.hpp b/cpp/include/raft/neighbors/hnsw_types.hpp index 645a0903b7..f90de6f01b 100644 --- a/cpp/include/raft/neighbors/hnsw_types.hpp +++ b/cpp/include/raft/neighbors/hnsw_types.hpp @@ -62,6 +62,11 @@ struct index : ann::index { auto metric() const -> raft::distance::DistanceType { return metric_; } + /** + @brief Set ef for search + */ + virtual void set_ef(int ef) const; + private: int dim_; raft::distance::DistanceType metric_; diff --git a/cpp/src/raft_runtime/neighbors/hnsw.cpp b/cpp/src/raft_runtime/neighbors/hnsw.cpp index 6eb770abd6..5356e708d2 100644 --- a/cpp/src/raft_runtime/neighbors/hnsw.cpp +++ b/cpp/src/raft_runtime/neighbors/hnsw.cpp @@ -21,6 +21,8 @@ #include #include +#include +#include namespace raft::neighbors::hnsw { #define RAFT_INST_HNSW(T) \ @@ -28,7 +30,11 @@ namespace raft::neighbors::hnsw { std::unique_ptr> from_cagra( \ raft::resources const& res, raft::neighbors::cagra::index cagra_index) \ { \ - std::string filepath = "/tmp/cagra_index.bin"; \ + std::random_device dev; \ + std::mt19937 rng(dev()); \ + std::uniform_int_distribution dist(0); \ + auto uuid = std::to_string(dist(rng)); \ + std::string filepath = "/tmp/" + uuid + ".bin"; \ raft::runtime::neighbors::cagra::serialize_to_hnswlib(res, filepath, cagra_index); \ auto hnsw_index = raft::runtime::neighbors::hnsw::deserialize_file( \ res, filepath, cagra_index.dim(), cagra_index.metric()); \ diff --git a/python/pylibraft/pylibraft/neighbors/hnsw.pyx b/python/pylibraft/pylibraft/neighbors/hnsw.pyx index aa589ffb65..e6f2d69eb8 100644 --- a/python/pylibraft/pylibraft/neighbors/hnsw.pyx +++ b/python/pylibraft/pylibraft/neighbors/hnsw.pyx @@ -52,6 +52,7 @@ from pylibraft.common.mdspan cimport ( from pylibraft.neighbors.common cimport _get_metric_string import os +import uuid import numpy as np @@ -292,7 +293,7 @@ def from_cagra(Index index, handle=None): Returns an hnswlib base-layer-only index from a CAGRA index. NOTE: This method uses the filesystem to write the CAGRA index in - `/tmp/cagra_index.bin` before reading it as an hnswlib index, + `/tmp/.bin` before reading it as an hnswlib index, then deleting the temporary file. Saving / loading the index is experimental. The serialization format is @@ -320,7 +321,8 @@ def from_cagra(Index index, handle=None): >>> # Serialize the CAGRA index to hnswlib base layer only index format >>> hnsw_index = hnsw.from_cagra(index, handle=handle) """ - filename = "/tmp/cagra_index.bin" + uuid_num = uuid.uuid4() + filename = f"/tmp/{uuid_num}.bin" save(filename, index, handle=handle) hnsw_index = load(filename, index.dim, np.dtype(index.active_index_type), _get_metric_string(index.metric), handle=handle) From ad4e5433f9309a7e20bd2154607cd8bda173204b Mon Sep 17 00:00:00 2001 From: Yinzuo Jiang Date: Thu, 27 Jun 2024 23:20:59 +0800 Subject: [PATCH 2/3] Fix compilation error when _CLK_BREAKDOWN is defined in cagra. (#2350) PR #1740 forgot to rename `BLOCK_SIZE` in `#ifdef _CLK_BREAKDOWN` blocks. The use of `RAFT_LOG_DEBUG` in kernel function results in compilation errors, replace it with `printf`. Also remove an unused function in search_single_cta_kernel-inl.cuh After merging: - [x] port to cuVS https://github.com/rapidsai/cuvs/pull/202 Authors: - Yinzuo Jiang (https://github.com/jiangyinzuo) - Tamas Bela Feher (https://github.com/tfeher) - Corey J. Nolet (https://github.com/cjnolet) Approvers: - tsuki (https://github.com/enp1s0) - Tamas Bela Feher (https://github.com/tfeher) URL: https://github.com/rapidsai/raft/pull/2350 --- .../cagra/search_multi_cta_kernel-inl.cuh | 12 ++++++++---- .../cagra/search_single_cta_kernel-inl.cuh | 18 +++++++----------- 2 files changed, 15 insertions(+), 15 deletions(-) diff --git a/cpp/include/raft/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh b/cpp/include/raft/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh index cfbb1e100c..16bb555aa4 100644 --- a/cpp/include/raft/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh +++ b/cpp/include/raft/neighbors/detail/cagra/search_multi_cta_kernel-inl.cuh @@ -36,6 +36,7 @@ #include #include +#include #include #include #include @@ -209,7 +210,7 @@ __launch_bounds__(1024, 1) RAFT_KERNEL search_kernel( #if 0 /* debug */ - for (unsigned i = threadIdx.x; i < result_buffer_size_32; i += BLOCK_SIZE) { + for (unsigned i = threadIdx.x; i < result_buffer_size_32; i += blockDim.x) { result_indices_buffer[i] = utils::get_max_value(); result_distances_buffer[i] = utils::get_max_value(); } @@ -351,16 +352,19 @@ __launch_bounds__(1024, 1) RAFT_KERNEL search_kernel( } #ifdef _CLK_BREAKDOWN - if ((threadIdx.x == 0 || threadIdx.x == BLOCK_SIZE - 1) && (blockIdx.x == 0) && + if ((threadIdx.x == 0 || threadIdx.x == blockDim.x - 1) && (blockIdx.x == 0) && ((query_id * 3) % gridDim.y < 3)) { - RAFT_LOG_DEBUG( + printf( + "%s:%d " "query, %d, thread, %d" - ", init, %d" + ", init, %lu" ", 1st_distance, %lu" ", topk, %lu" ", pickup_parents, %lu" ", distance, %lu" "\n", + __FILE__, + __LINE__, query_id, threadIdx.x, clk_init, diff --git a/cpp/include/raft/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh b/cpp/include/raft/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh index e8104bd6f6..232dcb782a 100644 --- a/cpp/include/raft/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh +++ b/cpp/include/raft/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh @@ -40,6 +40,7 @@ #include #include #include +#include #include #include #include @@ -448,14 +449,6 @@ __device__ inline void hashmap_restore(INDEX_T* const hashmap_ptr, } } -template -__device__ inline void set_value_device(T* const ptr, const T fill, const std::uint32_t count) -{ - for (std::uint32_t i = threadIdx.x; i < count; i += BLOCK_SIZE) { - ptr[i] = fill; - } -} - // One query one thread block template Date: Thu, 27 Jun 2024 12:03:46 -0400 Subject: [PATCH 3/3] Fix 0 recall issue in `raft_cagra_hnswlib` ANN benchmark (#2369) `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: https://github.com/rapidsai/raft/pull/2369 --- cpp/bench/ann/src/raft/raft_cagra_hnswlib_wrapper.h | 2 +- cpp/bench/ann/src/raft/raft_cagra_wrapper.h | 12 +++++++++--- 2 files changed, 10 insertions(+), 4 deletions(-) diff --git a/cpp/bench/ann/src/raft/raft_cagra_hnswlib_wrapper.h b/cpp/bench/ann/src/raft/raft_cagra_hnswlib_wrapper.h index 1c4b847d1a..1d2a1076ab 100644 --- a/cpp/bench/ann/src/raft/raft_cagra_hnswlib_wrapper.h +++ b/cpp/bench/ann/src/raft/raft_cagra_hnswlib_wrapper.h @@ -31,7 +31,7 @@ class RaftCagraHnswlib : public ANN, public AnnGPU { RaftCagraHnswlib(Metric metric, int dim, const BuildParam& param, int concurrent_searches = 1) : ANN(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::BuildParam{50, 100}} { diff --git a/cpp/bench/ann/src/raft/raft_cagra_wrapper.h b/cpp/bench/ann/src/raft/raft_cagra_wrapper.h index 0b892dec35..b03f875a8e 100644 --- a/cpp/bench/ann/src/raft/raft_cagra_wrapper.h +++ b/cpp/bench/ann/src/raft/raft_cagra_wrapper.h @@ -72,11 +72,16 @@ class RaftCagra : public ANN, public AnnGPU { std::optional 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(metric, dim), index_params_(param), dimension_(dim), need_dataset_update_(true), + shall_include_dataset_(shall_include_dataset), dataset_(std::make_shared>( std::move(make_device_matrix(handle_, 0, 0)))), graph_(std::make_shared>( @@ -135,6 +140,7 @@ class RaftCagra : public ANN, 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> index_; int dimension_; @@ -161,7 +167,7 @@ void RaftCagra::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>( std::move(raft::neighbors::cagra::detail::build(handle_, @@ -171,7 +177,7 @@ void RaftCagra::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)