Enabled C++ memory-map usage in PECOS-HNSW

pecos/core/ann/quantizer_impl/common.hpp

@@ -15,79 +15,54 @@
 #include <limits>
 #include <vector>
 #include "utils/clustering.hpp"
+#include "utils/mmap_util.hpp"
 
 namespace pecos {
 
 namespace ann {
 
-
     struct ProductQuantizer4BitsBase {
         // num_of_local_centroids denotes number of cluster centers used in quantization
         // In 4 Bit case, it's a fixed to be 16
-        const size_t num_of_local_centroids = 16;
+        const index_type num_of_local_centroids = 16;
         // num_local_codebooks denotes number of local codebooks we have or in other words,
         // number of subspace we have in Product Quantization.
         // Supposedly, num_local_codebooks * local_dimension equals dimension of original data vector
         index_type num_local_codebooks;
         // local dimension denotes the dimensionality of subspace in Product Quantization
         int local_dimension;
-        std::vector<float> global_centroid;
-        std::vector<float> local_codebooks;
-        std::vector<float> original_local_codebooks;
-
-        inline void save(FILE* fp) const {
-            pecos::file_util::fput_multiple<index_type>(&num_local_codebooks, 1, fp);
-            pecos::file_util::fput_multiple<int>(&local_dimension, 1, fp);
-            size_t sz = global_centroid.size();
-            pecos::file_util::fput_multiple<size_t>(&sz, 1, fp);
-            if (sz) {
-                pecos::file_util::fput_multiple<float>(&global_centroid[0], sz, fp);
-            }
-            sz = original_local_codebooks.size();
-            pecos::file_util::fput_multiple<size_t>(&sz, 1, fp);
-            if (sz) {
-                pecos::file_util::fput_multiple<float>(&original_local_codebooks[0], sz, fp);
-            }
-            sz = local_codebooks.size();
-            pecos::file_util::fput_multiple<size_t>(&sz, 1, fp);
-            if (sz) {
-                pecos::file_util::fput_multiple<float>(&local_codebooks[0], sz, fp);
-            }
+        mmap_util::MmapableVector<float> global_centroid;
+        mmap_util::MmapableVector<float> local_codebooks;
+        mmap_util::MmapableVector<float> original_local_codebooks;
+
+        inline void save(mmap_util::MmapStore& mmap_s) const {
+            mmap_s.fput_one<index_type>(this->num_local_codebooks);
+            mmap_s.fput_one<int>(this->local_dimension);
+            this->global_centroid.save_to_mmap_store(mmap_s);
+            this->local_codebooks.save_to_mmap_store(mmap_s);
+            this->original_local_codebooks.save_to_mmap_store(mmap_s);
         }
 
-        inline void load(FILE* fp) {
-            pecos::file_util::fget_multiple<index_type>(&num_local_codebooks, 1, fp);
-            pecos::file_util::fget_multiple<int>(&local_dimension, 1, fp);
-            size_t sz = 0;
-            pecos::file_util::fget_multiple<size_t>(&sz, 1, fp);
-            global_centroid.resize(sz);
-            if (sz) {
-                pecos::file_util::fget_multiple<float>(&global_centroid[0], sz, fp);
-            }
-            pecos::file_util::fget_multiple<size_t>(&sz, 1, fp);
-            original_local_codebooks.resize(sz);
-            if (sz) {
-                pecos::file_util::fget_multiple<float>(&original_local_codebooks[0], sz, fp);
-            }
-            pecos::file_util::fget_multiple<size_t>(&sz, 1, fp);
-            local_codebooks.resize(sz);
-            if (sz) {
-                pecos::file_util::fget_multiple<float>(&local_codebooks[0], sz, fp);
-            }
+        inline void load(mmap_util::MmapStore& mmap_s) {
+            this->num_local_codebooks = mmap_s.fget_one<index_type>();
+            this->local_dimension = mmap_s.fget_one<int>();
+            this->global_centroid.load_from_mmap_store(mmap_s);
+            this->local_codebooks.load_from_mmap_store(mmap_s);
+            this->original_local_codebooks.load_from_mmap_store(mmap_s);
         }
 
         inline void pack_codebook_for_inference_default() {
             local_codebooks = original_local_codebooks;
         }
-    
+
         inline void pad_parameters_default(index_type& max_degree, size_t& code_dimension) {}
-    
+
         inline void approximate_neighbor_group_distance_default(size_t neighbor_size, float* ds, const char* neighbor_codes, uint8_t* lut_ptr, float scale, float bias) const {
             index_type num_groups = neighbor_size % 16 == 0 ? neighbor_size / 16 : neighbor_size / 16 + 1;
-    
+
             std::vector<uint32_t> d(num_of_local_centroids);
             int ptr = 0;
-    
+
             const uint8_t *localID = reinterpret_cast<const uint8_t*>(neighbor_codes);
             for (index_type iters = 0; iters < num_groups; iters++) {
                 memset(d.data(), 0, sizeof(uint32_t) * num_of_local_centroids);
@@ -103,7 +78,7 @@ namespace ann {
                         }
                         d[k] += *(local_lut_ptr + obj);
                     }
-    
+
                     local_lut_ptr += num_of_local_centroids;
                 }
                 for (size_t k = 0; k < num_of_local_centroids; k++) {
@@ -112,7 +87,7 @@ namespace ann {
                 ptr += num_of_local_centroids;
             }
         }
-    
+
         inline void setup_lut_default(float* query, uint8_t* lut_ptr, float& scale, float& bias) const {
             float min = std::numeric_limits<float>::max();
             float max = std::numeric_limits<float>::min();
@@ -134,7 +109,7 @@ namespace ann {
                     min = std::min(min, tmp_v);
                 }
             }
-    
+
             bias = min;
             scale = (max - min) / 255.0;
             // second iteration to calculate quantized distnace and put it into lut
@@ -236,10 +211,7 @@ namespace ann {
                     &original_local_codebooks[m * num_of_local_centroids * local_dimension], threads);
             }
         }
-
     };
 
-
 }  // end of namespace ann
 }  // end of namespace pecos
-

test/pecos/ann/test_hnsw.py

@@ -82,6 +82,15 @@ def calc_recall(Y_true, Y_pred):
     Y_true = np.argsort(Y_true)[:, :top_k]
 
     # test dense features
+    #train_params = HNSW.TrainParams(M=36, efC=90, metric_type="ip", threads=1)
+    #pred_params = HNSW.PredParams(efS=80, topk=10, threads=1)
+    #model = HNSW.train(
+    #    X_trn,
+    #    train_params=train_params,
+    #    pred_params=pred_params,
+    #)
+    #model.save(dense_model_folder)
+    #del model
     model = HNSW.load(dense_model_folder)
     searchers = model.searchers_create(num_searcher_online)
     pred_params = model.get_pred_params()
@@ -99,6 +108,13 @@ def calc_recall(Y_true, Y_pred):
     # test csr features, we just reuse the Y_true since data are the same
     X_trn = smat.csr_matrix(X_trn).astype(np.float32)
     X_tst = smat.csr_matrix(X_tst).astype(np.float32)
+    #model = HNSW.train(
+    #    X_trn,
+    #    train_params=train_params,
+    #    pred_params=pred_params,
+    #)
+    #model.save(sparse_model_folder)
+    #del model
     model = HNSW.load(sparse_model_folder)
     searchers = model.searchers_create(num_searcher_online)
     pred_params = model.get_pred_params()

test/tst-data/ann/hnsw-model-dense/c_model/config.json

@@ -1,12 +1,12 @@
 {
     "hnsw_t": "pecos::ann::HNSW<float, pecos::ann::FeatVecDenseIPSimd<float>>",
-    "version": "v1.0",
     "train_params": {
-        "efC": 100,
-        "init_node": 50,
-        "maxM": 24,
-        "maxM0": 48,
+        "efC": 90,
+        "init_node": 66,
+        "maxM": 36,
+        "maxM0": 72,
         "max_level": 1,
         "num_node": 90
-    }
-}
+    },
+    "version": "v2.0"
+}

test/tst-data/ann/hnsw-model-dense/param.json

@@ -11,8 +11,8 @@
   "__meta__": {
    "class_fullname": "pecos.ann.hnsw.model###HNSW.PredParams"
   },
-  "efS": 100,
+  "efS": 80,
   "topk": 10,
   "threads": 1
  }
-}
+}

test/tst-data/ann/hnsw-model-sparse/c_model/config.json

@@ -1,12 +1,12 @@
 {
     "hnsw_t": "pecos::ann::HNSW<float, pecos::ann::FeatVecSparseIPSimd<uint32_t, float>>",
-    "version": "v1.0",
     "train_params": {
-        "efC": 100,
-        "init_node": 50,
-        "maxM": 24,
-        "maxM0": 48,
+        "efC": 90,
+        "init_node": 66,
+        "maxM": 36,
+        "maxM0": 72,
         "max_level": 1,
         "num_node": 90
-    }
-}
+    },
+    "version": "v2.0"
+}

test/tst-data/ann/hnsw-model-sparse/param.json

@@ -15,4 +15,4 @@
   "topk": 10,
   "threads": 1
  }
-}
+}