Skip to content

Commit

Permalink
Enabled C++ memory-map usage in PECOS-HNSW
Browse files Browse the repository at this point in the history
  • Loading branch information
Wei-Cheng Chang committed Feb 23, 2023
1 parent cef885f commit 706ddc3
Show file tree
Hide file tree
Showing 8 changed files with 160 additions and 231 deletions.
297 changes: 127 additions & 170 deletions pecos/core/ann/hnsw.hpp

Large diffs are not rendered by default.

78 changes: 25 additions & 53 deletions pecos/core/ann/quantizer_impl/common.hpp
Original file line number Diff line number Diff line change
Expand Up @@ -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);
Expand All @@ -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++) {
Expand All @@ -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();
Expand All @@ -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
Expand Down Expand Up @@ -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

6 changes: 3 additions & 3 deletions test/tst-data/ann/hnsw-model-dense/c_model/config.json
Original file line number Diff line number Diff line change
@@ -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,
"max_level": 1,
"num_node": 90
}
}
},
"version": "v2.0"
}
Binary file not shown.
2 changes: 1 addition & 1 deletion test/tst-data/ann/hnsw-model-dense/param.json
Original file line number Diff line number Diff line change
Expand Up @@ -15,4 +15,4 @@
"topk": 10,
"threads": 1
}
}
}
6 changes: 3 additions & 3 deletions test/tst-data/ann/hnsw-model-sparse/c_model/config.json
Original file line number Diff line number Diff line change
@@ -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,
"max_level": 1,
"num_node": 90
}
}
},
"version": "v2.0"
}
Binary file not shown.
2 changes: 1 addition & 1 deletion test/tst-data/ann/hnsw-model-sparse/param.json
Original file line number Diff line number Diff line change
Expand Up @@ -15,4 +15,4 @@
"topk": 10,
"threads": 1
}
}
}

0 comments on commit 706ddc3

Please sign in to comment.