update

Author: hhy3

glass/memory.hpp

@@ -3,25 +3,21 @@
 #include <sys/mman.h>
 
 #include <cstdint>
+#include <cstdio>
 #include <cstdlib>
 #include <cstring>
 
-#include "glass/common.hpp"
-
 namespace glass {
 
 constexpr size_t size_64B = 64;
 constexpr size_t size_2M = 2 * 1024 * 1024;
 constexpr size_t size_1G = 1 * 1024 * 1024 * 1024;
 
-template <size_t alignment>
-inline void *align_alloc_memory(size_t nbytes, bool set = true, uint8_t x = 0) {
+inline void *align_alloc_memory(size_t alignment, size_t nbytes, bool set = true, uint8_t x = 0) {
     size_t len = (nbytes + alignment - 1) / alignment * alignment;
-    if (alignment == size_1G) {
-        printf("Allocating %.2fG memory for %.2fG data\n", double(len) / size_1G, double(nbytes) / size_1G);
-    }
     auto p = std::aligned_alloc(alignment, len);
-    if constexpr (alignment >= size_2M) {
+    if (alignment >= size_2M) {
+        printf("Allocate %.2fMB for %.2fMB data\n", double(len) / 1024 / 1024, double(nbytes) / 1024 / 1024);
         madvise(p, len, MADV_HUGEPAGE);
     }
     if (set) {
@@ -31,12 +27,12 @@ inline void *align_alloc_memory(size_t nbytes, bool set = true, uint8_t x = 0) {
 }
 
 inline void *align_alloc(size_t nbytes, bool set = true, uint8_t x = 0) {
-    if (nbytes >= size_1G) {
-        return align_alloc_memory<size_1G>(nbytes, set, x);
+    if (nbytes >= size_1G / 2) {
+        return align_alloc_memory(size_1G, nbytes, set, x);
     } else if (nbytes >= size_2M) {
-        return align_alloc_memory<size_2M>(nbytes, set, x);
+        return align_alloc_memory(size_2M, nbytes, set, x);
     } else {
-        return align_alloc_memory<size_64B>(nbytes, set, x);
+        return align_alloc_memory(size_64B, nbytes, set, x);
     }
 }
 

glass/quant/computer.hpp

@@ -29,11 +29,11 @@ concept SymComputerConcept = ComputerBaseConcept<Computer> && requires(Computer
 
 template <StorageConcept Storage>
 struct Computer {
-    const Storage &tensor;
+    const Storage &storage;
 
-    explicit Computer(const Tensor &tensor) : tensor(tensor) {}
+    explicit Computer(const Storage &storage) : storage(storage) {}
 
-    void prefetch(int32_t u, int32_t lines) const { tensor.prefetch(u, lines); }
+    void prefetch(int32_t u, int32_t lines) const { storage.prefetch(u, lines); }
 };
 
 struct MemCpyTag {};
@@ -51,12 +51,12 @@ struct ComputerImpl : Computer<Storage> {
     mutable int64_t dist_cmps_{};
     mutable int64_t mem_read_bytes_{};
 
-    ComputerImpl(const Storage &tensor, const S *query, const auto &encoder, Args &&...args)
-        : Computer<Storage>(tensor), args(std::forward<Args>(args)...) {
+    ComputerImpl(const Storage &storage, const S *query, const auto &encoder, Args &&...args)
+        : Computer<Storage>(storage), args(std::forward<Args>(args)...) {
         if constexpr (std::is_same_v<std::decay_t<decltype(encoder)>, MemCpyTag>) {
             static_assert(std::is_same_v<S, X>);
-            q = (X *)align_alloc(this->tensor.dim_align() * sizeof(X));
-            memcpy(q, query, this->tensor.dim() * sizeof(X));
+            q = (X *)align_alloc(this->storage.dim_align() * sizeof(X));
+            memcpy(q, query, this->storage.dim() * sizeof(X));
         } else {
             encoder((const S *)query, q);
         }
@@ -66,11 +66,10 @@ struct ComputerImpl : Computer<Storage> {
 
     GLASS_INLINE dist_type operator()(const Y *p) const {
         dist_cmps_++;
-        mem_read_bytes_ += this->tensor.code_size();
-        return std::apply([&](auto &&...args) { return dist_func(q, p, this->tensor.dim_align(), args...); }, args);
+        return std::apply([&](auto &&...args) { return dist_func(q, p, this->storage.dim_align(), args...); }, args);
     }
 
-    GLASS_INLINE dist_type operator()(int32_t u) const { return operator()((const Y *)this->tensor.get(u)); }
+    GLASS_INLINE dist_type operator()(int32_t u) const { return operator()((const Y *)this->storage.get(u)); }
 
     GLASS_INLINE size_t dist_cmps() const { return dist_cmps_; }
 
@@ -86,15 +85,15 @@ struct SymComputerImpl : Computer<Storage> {
 
     std::tuple<Args...> args;
 
-    SymComputerImpl(const Storage &tensor, Args &&...args)
-        : Computer<Storage>(tensor), args(std::forward<Args>(args)...) {}
+    SymComputerImpl(const Storage &storage, Args &&...args)
+        : Computer<Storage>(storage), args(std::forward<Args>(args)...) {}
 
     GLASS_INLINE dist_type operator()(const X *x, const X *y) const {
-        return std::apply([&](auto &&...args) { return dist_func(x, y, this->tensor.dim_align(), args...); }, args);
+        return std::apply([&](auto &&...args) { return dist_func(x, y, this->storage.dim_align(), args...); }, args);
     }
 
     GLASS_INLINE dist_type operator()(int32_t u, int32_t v) const {
-        return operator()((const X *)this->tensor.get(u), (const X *)this->tensor.get(v));
+        return operator()((const X *)this->storage.get(u), (const X *)this->storage.get(v));
     }
 };
 

glass/searcher/graph_searcher.hpp

@@ -39,7 +39,7 @@ struct GraphSearcher : public SearcherBase {
     int32_t sample_points_num;
     std::vector<float> optimize_queries;
 
-    mutable std::vector<LinearPool<typename Quant::ComputerType::dist_type, Bitset<>>> pools;
+    mutable std::vector<LinearPool<typename Quant::ComputerType::dist_type, TwoLevelBitset<>>> pools;
 
     GraphSearcher(Graph<int32_t> g)
         : graph(std::move(g)), graph_po(graph.K / 16), pools(std::thread::hardware_concurrency()) {
@@ -151,7 +151,6 @@ struct GraphSearcher : public SearcherBase {
             latencies.resize(nq);
         }
         std::atomic<int64_t> total_dist_cmps{0};
-        std::atomic<int64_t> total_mem_read_bytes{0};
 #pragma omp parallel for schedule(dynamic)
         for (int32_t i = 0; i < nq; ++i) {
             std::chrono::high_resolution_clock::time_point start;
@@ -170,15 +169,15 @@ struct GraphSearcher : public SearcherBase {
             if (stats_enabled) {
                 auto end = std::chrono::high_resolution_clock::now();
                 latencies[i] = std::chrono::duration<float, std::milli>(end - start).count();
-                total_dist_cmps.fetch_add(computer.dist_cmps());
-                total_mem_read_bytes.fetch_add(computer.mem_read_bytes());
+                total_dist_cmps.fetch_add(computer.dist_cmps(), std::memory_order_relaxed);
             }
         }
         if (stats_enabled) {
             std::sort(latencies.begin(), latencies.end());
             stats.p99_latency_ms = latencies.empty() ? 0.0f : latencies[static_cast<size_t>(0.99 * nq)];
-            stats.avg_dist_comps = (double)total_dist_cmps.load() / nq;
-            stats.mem_read_bytes = total_mem_read_bytes.load();
+            size_t total_dist = total_dist_cmps.load(std::memory_order_relaxed);
+            stats.avg_dist_comps = (double)total_dist / nq;
+            stats.mem_read_bytes = total_dist * quant.code_size();
         }
     }