Engineering cleanups

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
-    $<$<AND:$<CONFIG:Release>,$<COMPILE_LANG_AND_ID:CXX,Clang,GNU>>:-Ofast>
+    $<$<AND:$<CONFIG:Release>,$<COMPILE_LANG_AND_ID:CXX,Clang,GNU>>:-O3>
     $<$<AND:$<BOOL:${RABITQ_ENABLE_NATIVE_OPTIMIZATION}>,$<COMPILE_LANG_AND_ID:CXX,Clang,GNU>>:-march=native>
 )
 target_link_libraries(rabitq_compile_options INTERFACE OpenMP::OpenMP_CXX)

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<float> qg;
@@ -97,8 +99,9 @@ qg.load("./qg_example.index"); // load pre-constructed index
 size_t ef = 100;
 size_t topk = 10;
 std::vector<PID> results(topk); // result buffer
+std::vector<float> dists(topk); // distance buffer
 std::vector<float> 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());
 ```

include/rabitqlib/index/hnsw/hnsw.hpp

@@ -192,7 +192,7 @@ class HierarchicalNSW {
 
     float (*ip_func_)(const float*, const uint8_t*, size_t);
 
-    Rotator<float>* rotator_ = nullptr;
+    std::unique_ptr<Rotator<float>> 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<float> : euclidean_sqr<float>) {
     max_elements_ = max_elements;
     dim_ = dim;
-    rotator_ = choose_rotator<float>(
+    rotator_.reset(choose_rotator<float>(
         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<VisitedListPool>(1, max_elements_);
 
-    rotator_ = choose_rotator<float>(
+    rotator_.reset(choose_rotator<float>(
         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);

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<float>* rotator_ = nullptr;              // Data Rotator
+    std::unique_ptr<Rotator<float>> rotator_;        // Data Rotator
     std::vector<Cluster> 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<float>(dim, type, round_up_to_multiple(dim_, 64));
+    rotator_.reset(choose_rotator<float>(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<char*>(&type_), sizeof(type_));
     input.read(reinterpret_cast<char*>(&metric_type_), sizeof(metric_type_));
 
-    delete rotator_;
-    rotator_ = choose_rotator<float>(dim_, type_, round_up_to_multiple(dim_, 64));
+    rotator_.reset(choose_rotator<float>(dim_, type_, round_up_to_multiple(dim_, 64)));
     padded_dim_ = rotator_->size();
 
     /* Load number of vectors of each cluster */

include/rabitqlib/index/symqg/qg.hpp

@@ -9,6 +9,7 @@
 #include <cstring>
 #include <fstream>
 #include <iostream>
+#include <memory>
 #include <ostream>
 #include <stdexcept>
 #include <vector>
@@ -51,8 +52,8 @@ class QuantizedGraph {
             char,
             1 << 22,
             true>>
-        data_;                       // vectors + graph + quantization codes + factors
-    Rotator<T>* rotator_ = nullptr;  // data rotator
+        data_;                             // vectors + graph + quantization codes + factors
+    std::unique_ptr<Rotator<T>> rotator_;  // data rotator
     std::unique_ptr<VisitedListPool> 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<T>::validate_configuration() const {
     }
 }
 
-template <typename T>
-inline QuantizedGraph<T>::~QuantizedGraph() {
-    delete this->rotator_;
-}
-
 template <typename T>
 inline void QuantizedGraph<T>::copy_vectors(const T* data) {
 #pragma omp parallel for schedule(dynamic)
@@ -270,52 +265,6 @@ inline void QuantizedGraph<T>::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 <typename T>
-inline void QuantizedGraph<T>::search(
-    const T* __restrict__ query, uint32_t k, uint32_t* __restrict__ results
-) {
-    std::vector<T> rotated_query(padded_dim_);
-    rotator_->rotate(query, rotated_query.data());
-
-    // init query
-    BatchQuery<T> q_obj(rotated_query.data(), padded_dim_);
-
-    buffer::SearchBuffer<T> search_pool(ef_);
-    // init search buffer
-    search_pool.insert(this->entry_point_, std::numeric_limits<T>::max());
-
-    buffer::SearchBuffer res_pool(k);  // result buffer
-    auto* vis = visited_list_pool_->get_free_vislist();
-
-    std::vector<T> 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 <typename T>
 inline void QuantizedGraph<T>::search(
     const T* __restrict__ query,
@@ -419,9 +368,9 @@ inline void QuantizedGraph<T>::update_results(
 // initialize const offsets & data array
 template <typename T>
 inline void QuantizedGraph<T>::initialize() {
-    delete rotator_;
-
-    rotator_ = choose_rotator<float>(dim_, rotator_type_, round_up_to_multiple(dim_, 64));
+    rotator_.reset(
+        choose_rotator<float>(dim_, rotator_type_, round_up_to_multiple(dim_, 64))
+    );
     padded_dim_ = rotator_->size();
 
     /* check size */

python_bindings/ivf_bindings.cpp

@@ -1,5 +1,6 @@
 #include <pybind11/stl.h>
 
+#include <algorithm>
 #include <limits>
 #include <memory>
 #include <string>
@@ -105,32 +106,26 @@ class IvfIndex {
             std::vector<ssize_t>{static_cast<ssize_t>(nq), static_cast<ssize_t>(k)};
         auto ids = py::array_t<rabitqlib::PID>(shape);
         auto dists = py::array_t<float>(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<float>::infinity()
+        );
 
         rabitqlib::ivf::parallel_for(
             0,
             nq,
             num_threads,
             [&](size_t idx, size_t /*threadId*/) {
-                std::vector<rabitqlib::PID> row_ids(k, rabitqlib::kPidMax);
-                std::vector<float> row_dists(k, std::numeric_limits<float>::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<ssize_t>(idx), static_cast<ssize_t>(j)) =
-                        row_ids[j];
-                    dists_buf(static_cast<ssize_t>(idx), static_cast<ssize_t>(j)) =
-                        row_dists[j];
-                }
             }
         );
 

python_bindings/symqg_bindings.cpp

@@ -63,28 +63,22 @@ class SymqgIndex {
             std::vector<ssize_t>{static_cast<ssize_t>(nq), static_cast<ssize_t>(k)};
         auto ids = py::array_t<rabitqlib::PID>(shape);
         auto dists = py::array_t<float>(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<rabitqlib::PID> row_ids(k, 0);
-                std::vector<float> row_dists(k, 0.0F);
                 index_->search(
                     query_array.data() + (idx * dim_),
                     static_cast<uint32_t>(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<ssize_t>(idx), static_cast<ssize_t>(j)) =
-                        row_ids[j];
-                    dists_buf(static_cast<ssize_t>(idx), static_cast<ssize_t>(j)) =
-                        row_dists[j];
-                }
             }
         );
 

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<PID> results(topk);
+            std::vector<float> dists(topk);
             for (size_t z = 0; z < nq; z++) {
                 stopw.reset();
-                qg.search(&query(static_cast<Eigen::Index>(z), 0), topk, results.data());
+                qg.search(
+                    &query(static_cast<Eigen::Index>(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++) {