diff --git a/CMakeLists.txt b/CMakeLists.txt index d70f2da..65029b6 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -21,7 +21,7 @@ find_package(OpenMP REQUIRED) add_library(rabitq_compile_options INTERFACE) target_compile_features(rabitq_compile_options INTERFACE cxx_std_17) target_compile_options(rabitq_compile_options INTERFACE - $<$,$>:-Ofast> + $<$,$>:-O3> $<$,$>:-march=native> ) target_link_libraries(rabitq_compile_options INTERFACE OpenMP::OpenMP_CXX) diff --git a/docs/docs/index/qg.md b/docs/docs/index/qg.md index f215cce..17ff2ef 100644 --- a/docs/docs/index/qg.md +++ b/docs/docs/index/qg.md @@ -81,13 +81,15 @@ element's neighbors, organized in FastScan batches of 32. Consequently, For querying, code is pretty simple. ```cpp void QuantizedGraph::search( - const T* __restrict__ query, - uint32_t k, - uint32_t* __restrict__ results); + const T* __restrict__ query, + uint32_t k, + uint32_t* __restrict__ results, + T* __restrict__ dists); ``` - **query**: Query vector. - **k**: Top-k. - **results**: Result buffer, size of k. +- **dists**: Distance buffer, size of k. Then we can use a pre-constructed index to search. ```cpp QuantizedGraph qg; @@ -97,8 +99,9 @@ qg.load("./qg_example.index"); // load pre-constructed index size_t ef = 100; size_t topk = 10; std::vector results(topk); // result buffer +std::vector dists(topk); // distance buffer std::vector query(cols); // populate with a query vector qg.set_ef(ef); // set search window size -qg.search(query.data(), topk, results.data()); +qg.search(query.data(), topk, results.data(), dists.data()); ``` diff --git a/include/rabitqlib/index/hnsw/hnsw.hpp b/include/rabitqlib/index/hnsw/hnsw.hpp index 6e5b0af..5e3a5c8 100644 --- a/include/rabitqlib/index/hnsw/hnsw.hpp +++ b/include/rabitqlib/index/hnsw/hnsw.hpp @@ -192,7 +192,7 @@ class HierarchicalNSW { float (*ip_func_)(const float*, const uint8_t*, size_t); - Rotator* rotator_ = nullptr; + std::unique_ptr> rotator_; quant::RabitqConfig query_config_; @@ -222,8 +222,7 @@ class HierarchicalNSW { free(centroids_memory_); - delete rotator_; - rotator_ = nullptr; + rotator_.reset(); } void set_ef(size_t ef) { ef_ = ef; } @@ -375,9 +374,9 @@ inline HierarchicalNSW::HierarchicalNSW( , raw_dist_func_((metric_type == METRIC_IP) ? dot_product_dis : euclidean_sqr) { max_elements_ = max_elements; dim_ = dim; - rotator_ = choose_rotator( + rotator_.reset(choose_rotator( dim, RotatorType::FhtKacRotator, round_up_to_multiple(dim_, 64) - ); + )); padded_dim_ = rotator_->size(); /* check size */ assert(padded_dim_ % 64 == 0); @@ -604,9 +603,9 @@ inline void HierarchicalNSW::load(const char* filename) { visited_list_pool_ = std::make_unique(1, max_elements_); - rotator_ = choose_rotator( + rotator_.reset(choose_rotator( dim_, RotatorType::FhtKacRotator, round_up_to_multiple(dim_, 64) - ); + )); if (rotator_->size() != padded_dim_) { std::cerr << "Bad padded_dim_ for rotator in hnsw.load()\n"; exit(1); diff --git a/include/rabitqlib/index/ivf/ivf.hpp b/include/rabitqlib/index/ivf/ivf.hpp index dc8600b..83ff26f 100644 --- a/include/rabitqlib/index/ivf/ivf.hpp +++ b/include/rabitqlib/index/ivf/ivf.hpp @@ -39,7 +39,7 @@ class IVF { size_t num_cluster_ = 0; // num of centroids (clusters) size_t ex_bits_ = 0; // total bits = ex_bits_ + 1 RotatorType type_ = RotatorType::FhtKacRotator; // type of rotator - Rotator* rotator_ = nullptr; // Data Rotator + std::unique_ptr> rotator_; // Data Rotator std::vector cluster_lst_; // List of clusters in ivf MetricType metric_type_ = rabitqlib::METRIC_L2; // metric type float (*ip_func_)(const float*, const uint8_t*, size_t) = nullptr; @@ -143,17 +143,14 @@ inline IVF::IVF( std::cerr.flush(); exit(1); }; - rotator_ = choose_rotator(dim, type, round_up_to_multiple(dim_, 64)); + rotator_.reset(choose_rotator(dim, type, round_up_to_multiple(dim_, 64))); padded_dim_ = rotator_->size(); /* check size */ assert(padded_dim_ % 64 == 0); assert(padded_dim_ >= dim_); } -inline IVF::~IVF() { - delete rotator_; - free_memory(); -} +inline IVF::~IVF() { free_memory(); } /** * @brief Construct clusters in IVF @@ -368,8 +365,7 @@ inline void IVF::load(const char* filename) { input.read(reinterpret_cast(&type_), sizeof(type_)); input.read(reinterpret_cast(&metric_type_), sizeof(metric_type_)); - delete rotator_; - rotator_ = choose_rotator(dim_, type_, round_up_to_multiple(dim_, 64)); + rotator_.reset(choose_rotator(dim_, type_, round_up_to_multiple(dim_, 64))); padded_dim_ = rotator_->size(); /* Load number of vectors of each cluster */ diff --git a/include/rabitqlib/index/symqg/qg.hpp b/include/rabitqlib/index/symqg/qg.hpp index 4cd9dc1..a3c26ce 100644 --- a/include/rabitqlib/index/symqg/qg.hpp +++ b/include/rabitqlib/index/symqg/qg.hpp @@ -9,6 +9,7 @@ #include #include #include +#include #include #include #include @@ -51,8 +52,8 @@ class QuantizedGraph { char, 1 << 22, true>> - data_; // vectors + graph + quantization codes + factors - Rotator* rotator_ = nullptr; // data rotator + data_; // vectors + graph + quantization codes + factors + std::unique_ptr> rotator_; // data rotator std::unique_ptr visited_list_pool_ = nullptr; // Position of different data in each row (RawData + QuantizationCodes + Factors + @@ -119,7 +120,7 @@ class QuantizedGraph { explicit QuantizedGraph() = default; - ~QuantizedGraph(); + ~QuantizedGraph() = default; [[nodiscard]] auto num_vertices() const { return this->num_points_; } @@ -140,7 +141,6 @@ class QuantizedGraph { void set_ef(size_t); /* search and copy results to KNN */ - void search(const T* __restrict__ query, uint32_t knn, uint32_t* __restrict__ results); void search( const T* __restrict__ query, uint32_t knn, @@ -183,11 +183,6 @@ inline void QuantizedGraph::validate_configuration() const { } } -template -inline QuantizedGraph::~QuantizedGraph() { - delete this->rotator_; -} - template inline void QuantizedGraph::copy_vectors(const T* data) { #pragma omp parallel for schedule(dynamic) @@ -270,52 +265,6 @@ inline void QuantizedGraph::set_ef(size_t cur_ef) { this->ef_ = cur_ef; } -/** - * @brief search on qg - * - * @param query unrotated query vector, dimension_ elements - * @param knn num of nearest neighbors - * @param results search result - */ -template -inline void QuantizedGraph::search( - const T* __restrict__ query, uint32_t k, uint32_t* __restrict__ results -) { - std::vector rotated_query(padded_dim_); - rotator_->rotate(query, rotated_query.data()); - - // init query - BatchQuery q_obj(rotated_query.data(), padded_dim_); - - buffer::SearchBuffer search_pool(ef_); - // init search buffer - search_pool.insert(this->entry_point_, std::numeric_limits::max()); - - buffer::SearchBuffer res_pool(k); // result buffer - auto* vis = visited_list_pool_->get_free_vislist(); - - std::vector est_dist(degree_bound_); // estimated distances - - while (search_pool.has_next()) { - PID cur_node = search_pool.pop(); - if (vis->get(cur_node)) { - continue; - } - vis->set(cur_node); - - q_obj.set_g_add(raw_dist_func_(query, get_vector(cur_node), dim_)); - - scan_neighbors( - q_obj, cur_node, est_dist.data(), search_pool, *vis, this->degree_bound_ - ); - res_pool.insert(cur_node, q_obj.g_add()); - } - - update_results(res_pool, *vis, query); - visited_list_pool_->release_vis_list(vis); - res_pool.copy_results(results); -} - template inline void QuantizedGraph::search( const T* __restrict__ query, @@ -419,9 +368,9 @@ inline void QuantizedGraph::update_results( // initialize const offsets & data array template inline void QuantizedGraph::initialize() { - delete rotator_; - - rotator_ = choose_rotator(dim_, rotator_type_, round_up_to_multiple(dim_, 64)); + rotator_.reset( + choose_rotator(dim_, rotator_type_, round_up_to_multiple(dim_, 64)) + ); padded_dim_ = rotator_->size(); /* check size */ diff --git a/python_bindings/ivf_bindings.cpp b/python_bindings/ivf_bindings.cpp index a9c2ea7..7ecc37b 100644 --- a/python_bindings/ivf_bindings.cpp +++ b/python_bindings/ivf_bindings.cpp @@ -1,5 +1,6 @@ #include +#include #include #include #include @@ -105,32 +106,26 @@ class IvfIndex { std::vector{static_cast(nq), static_cast(k)}; auto ids = py::array_t(shape); auto dists = py::array_t(shape); - auto ids_buf = ids.mutable_unchecked<2>(); - auto dists_buf = dists.mutable_unchecked<2>(); + auto* ids_data = ids.mutable_data(); + auto* dists_data = dists.mutable_data(); + std::fill(ids_data, ids_data + ids.size(), rabitqlib::kPidMax); + std::fill( + dists_data, dists_data + dists.size(), std::numeric_limits::infinity() + ); rabitqlib::ivf::parallel_for( 0, nq, num_threads, [&](size_t idx, size_t /*threadId*/) { - std::vector row_ids(k, rabitqlib::kPidMax); - std::vector row_dists(k, std::numeric_limits::infinity()); - index_->search( query_array.data() + (idx * dim_), k, nprobe, - row_ids.data(), - row_dists.data(), + ids_data + (idx * k), + dists_data + (idx * k), high_accuracy ); - - for (size_t j = 0; j < k; ++j) { - ids_buf(static_cast(idx), static_cast(j)) = - row_ids[j]; - dists_buf(static_cast(idx), static_cast(j)) = - row_dists[j]; - } } ); diff --git a/python_bindings/symqg_bindings.cpp b/python_bindings/symqg_bindings.cpp index 4008e9b..768bf2a 100644 --- a/python_bindings/symqg_bindings.cpp +++ b/python_bindings/symqg_bindings.cpp @@ -63,28 +63,22 @@ class SymqgIndex { std::vector{static_cast(nq), static_cast(k)}; auto ids = py::array_t(shape); auto dists = py::array_t(shape); - auto ids_buf = ids.mutable_unchecked<2>(); - auto dists_buf = dists.mutable_unchecked<2>(); + auto* ids_data = ids.mutable_data(); + auto* dists_data = dists.mutable_data(); + std::fill(ids_data, ids_data + ids.size(), 0); + std::fill(dists_data, dists_data + dists.size(), 0.0F); rabitqlib::ivf::parallel_for( 0, nq, num_threads, [&](size_t idx, size_t /*threadId*/) { - std::vector row_ids(k, 0); - std::vector row_dists(k, 0.0F); index_->search( query_array.data() + (idx * dim_), static_cast(k), - row_ids.data(), - row_dists.data() + ids_data + (idx * k), + dists_data + (idx * k) ); - for (size_t j = 0; j < k; ++j) { - ids_buf(static_cast(idx), static_cast(j)) = - row_ids[j]; - dists_buf(static_cast(idx), static_cast(j)) = - row_dists[j]; - } } ); diff --git a/sample/cpp/symqg_querying.cpp b/sample/cpp/symqg_querying.cpp index 5ae2047..4a97cdd 100644 --- a/sample/cpp/symqg_querying.cpp +++ b/sample/cpp/symqg_querying.cpp @@ -68,9 +68,15 @@ int main(int argc, char** argv) { float total_time = 0; qg.set_ef(ef); std::vector results(topk); + std::vector dists(topk); for (size_t z = 0; z < nq; z++) { stopw.reset(); - qg.search(&query(static_cast(z), 0), topk, results.data()); + qg.search( + &query(static_cast(z), 0), + topk, + results.data(), + dists.data() + ); total_time += stopw.get_elapsed_micro(); for (size_t y = 0; y < topk; y++) { for (size_t k = 0; k < topk; k++) {