deglib 0.2.3__tar.gz → 0.2.5__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (90) hide show
  1. {deglib-0.2.3/src/deglib.egg-info → deglib-0.2.5}/PKG-INFO +15 -19
  2. {deglib-0.2.3 → deglib-0.2.5}/README.md +14 -18
  3. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/distance/fp32_ip.h +20 -14
  4. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/distance/fp32_l2.h +20 -14
  5. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/distance/int8_l2.h +28 -18
  6. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/distance/uint8_l2.h +28 -18
  7. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/optimization/quantization/evp_quantize.h +80 -8
  8. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/optimization/quantization/scalar_quantize.h +161 -190
  9. deglib-0.2.5/lib/deglib/include/deglib/search/searcher.h +493 -0
  10. deglib-0.2.5/lib/deglib/include/deglib/search.h +253 -0
  11. {deglib-0.2.3 → deglib-0.2.5}/src/deg_cpp/deglib_cpp.cpp +23 -176
  12. {deglib-0.2.3 → deglib-0.2.5}/src/deglib/__init__.py +1 -1
  13. {deglib-0.2.3 → deglib-0.2.5}/src/deglib/search.py +31 -55
  14. {deglib-0.2.3 → deglib-0.2.5/src/deglib.egg-info}/PKG-INFO +15 -19
  15. {deglib-0.2.3 → deglib-0.2.5}/tests/test_builder.py +28 -0
  16. {deglib-0.2.3 → deglib-0.2.5}/tests/test_scalar_quantizer.py +22 -35
  17. {deglib-0.2.3 → deglib-0.2.5}/tests/test_searcher.py +33 -29
  18. deglib-0.2.3/lib/deglib/include/deglib/search/searcher.h +0 -454
  19. deglib-0.2.3/lib/deglib/include/deglib/search.h +0 -112
  20. {deglib-0.2.3 → deglib-0.2.5}/CMakeLists.txt +0 -0
  21. {deglib-0.2.3 → deglib-0.2.5}/LICENSE +0 -0
  22. {deglib-0.2.3 → deglib-0.2.5}/MANIFEST.in +0 -0
  23. {deglib-0.2.3 → deglib-0.2.5}/lib/CMakeLists.txt +0 -0
  24. {deglib-0.2.3 → deglib-0.2.5}/lib/bench/CMakeLists.txt +0 -0
  25. {deglib-0.2.3 → deglib-0.2.5}/lib/bench/include/benchmark.h +0 -0
  26. {deglib-0.2.3 → deglib-0.2.5}/lib/bench/include/build.h +0 -0
  27. {deglib-0.2.3 → deglib-0.2.5}/lib/bench/include/dataset.h +0 -0
  28. {deglib-0.2.3 → deglib-0.2.5}/lib/bench/include/file_io.h +0 -0
  29. {deglib-0.2.3 → deglib-0.2.5}/lib/bench/include/logging.h +0 -0
  30. {deglib-0.2.3 → deglib-0.2.5}/lib/bench/include/repository.h +0 -0
  31. {deglib-0.2.3 → deglib-0.2.5}/lib/bench/include/stats.h +0 -0
  32. {deglib-0.2.3 → deglib-0.2.5}/lib/bench/include/stopwatch.h +0 -0
  33. {deglib-0.2.3 → deglib-0.2.5}/lib/cmake_modules/DetectCPUFeatures.cmake +0 -0
  34. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/CMakeLists.txt +0 -0
  35. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/analysis.h +0 -0
  36. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/builder.h +0 -0
  37. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/concurrent.h +0 -0
  38. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/config.h +0 -0
  39. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/deglib.h +0 -0
  40. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/distance/evp_ip.h +0 -0
  41. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/distance/fp16.h +0 -0
  42. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/distance/fp16_ip.h +0 -0
  43. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/distance/fp16_l2.h +0 -0
  44. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/distance/fp32.h +0 -0
  45. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/distance/int8.h +0 -0
  46. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/distance/int8_ip.h +0 -0
  47. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/distance/residual_mode.h +0 -0
  48. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/distance/uint8.h +0 -0
  49. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/distance/uint8_ip.h +0 -0
  50. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/distances.h +0 -0
  51. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/filter.h +0 -0
  52. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/graph/dynamic_graph.h +0 -0
  53. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/graph/internal_graph.h +0 -0
  54. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/graph/intrusive_list.h +0 -0
  55. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/graph/mutable_graph.h +0 -0
  56. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/graph/readonly_graph.h +0 -0
  57. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/graph/sizebounded_graph.h +0 -0
  58. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/graph/visited_list_pool.h +0 -0
  59. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/graph.h +0 -0
  60. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/optimization/flas/fast_linear_assignment_sorter.h +0 -0
  61. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/optimization/flas/fast_linear_assignment_sorter_mt.h +0 -0
  62. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/optimization/flas/junker_volgenant_solver.h +0 -0
  63. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/optimization/pruning.h +0 -0
  64. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/optimization/transform.h +0 -0
  65. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/optimization.h +0 -0
  66. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/utils/cpu.h +0 -0
  67. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/utils/memory.h +0 -0
  68. {deglib-0.2.3 → deglib-0.2.5}/lib/deglib/include/deglib/utils/random.h +0 -0
  69. {deglib-0.2.3 → deglib-0.2.5}/lib/test/CMakeLists.txt +0 -0
  70. {deglib-0.2.3 → deglib-0.2.5}/lib/test/src/common/test_helpers.h +0 -0
  71. {deglib-0.2.3 → deglib-0.2.5}/pyproject.toml +0 -0
  72. {deglib-0.2.3 → deglib-0.2.5}/setup.cfg +0 -0
  73. {deglib-0.2.3 → deglib-0.2.5}/setup.py +0 -0
  74. {deglib-0.2.3 → deglib-0.2.5}/src/deglib/analysis.py +0 -0
  75. {deglib-0.2.3 → deglib-0.2.5}/src/deglib/builder.py +0 -0
  76. {deglib-0.2.3 → deglib-0.2.5}/src/deglib/cpu.py +0 -0
  77. {deglib-0.2.3 → deglib-0.2.5}/src/deglib/distances.py +0 -0
  78. {deglib-0.2.3 → deglib-0.2.5}/src/deglib/graph.py +0 -0
  79. {deglib-0.2.3 → deglib-0.2.5}/src/deglib/optimization.py +0 -0
  80. {deglib-0.2.3 → deglib-0.2.5}/src/deglib/utils.py +0 -0
  81. {deglib-0.2.3 → deglib-0.2.5}/src/deglib.egg-info/SOURCES.txt +0 -0
  82. {deglib-0.2.3 → deglib-0.2.5}/src/deglib.egg-info/dependency_links.txt +0 -0
  83. {deglib-0.2.3 → deglib-0.2.5}/src/deglib.egg-info/requires.txt +0 -0
  84. {deglib-0.2.3 → deglib-0.2.5}/src/deglib.egg-info/top_level.txt +0 -0
  85. {deglib-0.2.3 → deglib-0.2.5}/tests/test_analysis.py +0 -0
  86. {deglib-0.2.3 → deglib-0.2.5}/tests/test_filter.py +0 -0
  87. {deglib-0.2.3 → deglib-0.2.5}/tests/test_graph.py +0 -0
  88. {deglib-0.2.3 → deglib-0.2.5}/tests/test_presort.py +0 -0
  89. {deglib-0.2.3 → deglib-0.2.5}/tests/test_rerank.py +0 -0
  90. {deglib-0.2.3 → deglib-0.2.5}/tests/test_transform.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: deglib
3
- Version: 0.2.3
3
+ Version: 0.2.5
4
4
  Summary: Python bindings for the Dynamic Exploration Graph library by Nico Hezel
5
5
  Author: Nico Hezel
6
6
  Author-email: Bruno Schilling <bruno.schilling@htw-berlin.de>
@@ -216,17 +216,17 @@ explored_labels, distances = graph.explore(entry_external_label=105, k=10, eps=0
216
216
 
217
217
  ### Graph Optimization, Quantization & Reranking Pipeline
218
218
 
219
- A complete end-to-end pipeline demonstrating **FLAS pre-sorting**, **multithreaded graph building**, **RNG edge pruning**, **Int8 quantization**, **ReadOnlyGraph conversion**, and **exact FP32 candidate reranking**:
219
+ A complete end-to-end pipeline demonstrating **FLAS pre-sorting**, **multithreaded graph building**, **RNG edge pruning**, **Int8 quantization**, **ReadOnlyGraph conversion**, and deployment with **`Searcher`** for automatic query quantization and high-precision reranking:
220
220
 
221
221
  ```python
222
222
  import numpy as np
223
223
  import deglib
224
224
  from deglib.optimization import presort, prune_non_rng_edges, ScalarQuantizerInt8
225
- from deglib.search import rerank
225
+ from deglib.search import create_searcher
226
226
 
227
227
  num_vectors, dims = 10_000, 128
228
228
  data = np.random.randn(num_vectors, dims).astype(np.float32)
229
- query = np.random.randn(1, dims).astype(np.float32)
229
+ query = np.random.randn(dims).astype(np.float32)
230
230
 
231
231
  # 1. Pre-sort vectors using FLAS for improved memory locality and index construction speed
232
232
  perm = presort(data, metric=deglib.Metric.FP32_InnerProduct, callback="progress")
@@ -253,23 +253,19 @@ int8_space = deglib.FloatSpace.create(dims, deglib.Metric.Int8_InnerProduct)
253
253
  # 5. Convert mutable graph to a compact ReadOnlyGraph equipped with the quantized features
254
254
  readonly_graph = graph.to_readonly(feature_space=int8_space, custom_features=quant_int8_data)
255
255
 
256
- # 6. Quantize query using the EXACT SAME calibrated scale, then search
257
- quant_query = quantizer.quantize(query)
258
- candidate_indices = readonly_graph.search(quant_query, k=20, eps=0.1, return_distances=False)
259
-
260
- # 7. Exact distance reranking of top candidates on original float32 data
261
- fp32_space = deglib.FloatSpace.create(dims, deglib.Metric.FP32_InnerProduct)
262
- final_top_indices, distances = rerank(
263
- space=fp32_space,
264
- queries=query,
265
- candidate_indices=candidate_indices,
266
- base_vectors=sorted_data,
267
- k_top=10,
268
- return_distances=True,
256
+ # 6. Create high-performance Searcher equipped with base vectors for exact candidate reranking
257
+ rerank_space = deglib.FloatSpace.create(dims, deglib.Metric.FP32_InnerProduct)
258
+ searcher = create_searcher(
259
+ graph=readonly_graph,
260
+ quantizer=quantizer,
261
+ refine_space=rerank_space,
262
+ refine_data=sorted_data,
269
263
  )
270
264
 
271
- print("Top-10 nearest neighbor indices:", final_top_indices[0])
272
- print("Top-10 exact distances:", distances[0])
265
+ # 7. Query search (query is automatically quantized and top candidates are reranked)
266
+ indices, distances = searcher.search(query, k=10, eps=0.1, rerank_factor=1.5, return_distances=True)
267
+ print("Top-10 nearest neighbor indices:", indices)
268
+ print("Top-10 exact distances:", distances)
273
269
  ```
274
270
 
275
271
  ---
@@ -188,17 +188,17 @@ explored_labels, distances = graph.explore(entry_external_label=105, k=10, eps=0
188
188
 
189
189
  ### Graph Optimization, Quantization & Reranking Pipeline
190
190
 
191
- A complete end-to-end pipeline demonstrating **FLAS pre-sorting**, **multithreaded graph building**, **RNG edge pruning**, **Int8 quantization**, **ReadOnlyGraph conversion**, and **exact FP32 candidate reranking**:
191
+ A complete end-to-end pipeline demonstrating **FLAS pre-sorting**, **multithreaded graph building**, **RNG edge pruning**, **Int8 quantization**, **ReadOnlyGraph conversion**, and deployment with **`Searcher`** for automatic query quantization and high-precision reranking:
192
192
 
193
193
  ```python
194
194
  import numpy as np
195
195
  import deglib
196
196
  from deglib.optimization import presort, prune_non_rng_edges, ScalarQuantizerInt8
197
- from deglib.search import rerank
197
+ from deglib.search import create_searcher
198
198
 
199
199
  num_vectors, dims = 10_000, 128
200
200
  data = np.random.randn(num_vectors, dims).astype(np.float32)
201
- query = np.random.randn(1, dims).astype(np.float32)
201
+ query = np.random.randn(dims).astype(np.float32)
202
202
 
203
203
  # 1. Pre-sort vectors using FLAS for improved memory locality and index construction speed
204
204
  perm = presort(data, metric=deglib.Metric.FP32_InnerProduct, callback="progress")
@@ -225,23 +225,19 @@ int8_space = deglib.FloatSpace.create(dims, deglib.Metric.Int8_InnerProduct)
225
225
  # 5. Convert mutable graph to a compact ReadOnlyGraph equipped with the quantized features
226
226
  readonly_graph = graph.to_readonly(feature_space=int8_space, custom_features=quant_int8_data)
227
227
 
228
- # 6. Quantize query using the EXACT SAME calibrated scale, then search
229
- quant_query = quantizer.quantize(query)
230
- candidate_indices = readonly_graph.search(quant_query, k=20, eps=0.1, return_distances=False)
231
-
232
- # 7. Exact distance reranking of top candidates on original float32 data
233
- fp32_space = deglib.FloatSpace.create(dims, deglib.Metric.FP32_InnerProduct)
234
- final_top_indices, distances = rerank(
235
- space=fp32_space,
236
- queries=query,
237
- candidate_indices=candidate_indices,
238
- base_vectors=sorted_data,
239
- k_top=10,
240
- return_distances=True,
228
+ # 6. Create high-performance Searcher equipped with base vectors for exact candidate reranking
229
+ rerank_space = deglib.FloatSpace.create(dims, deglib.Metric.FP32_InnerProduct)
230
+ searcher = create_searcher(
231
+ graph=readonly_graph,
232
+ quantizer=quantizer,
233
+ refine_space=rerank_space,
234
+ refine_data=sorted_data,
241
235
  )
242
236
 
243
- print("Top-10 nearest neighbor indices:", final_top_indices[0])
244
- print("Top-10 exact distances:", distances[0])
237
+ # 7. Query search (query is automatically quantized and top candidates are reranked)
238
+ indices, distances = searcher.search(query, k=10, eps=0.1, rerank_factor=1.5, return_distances=True)
239
+ print("Top-10 nearest neighbor indices:", indices)
240
+ print("Top-10 exact distances:", distances)
245
241
  ```
246
242
 
247
243
  ---
@@ -143,17 +143,20 @@ class InnerProductFloat_AVX512 {
143
143
  }
144
144
 
145
145
  if constexpr (HasDualSimd) {
146
- const size_t nc16 = dim / 16;
147
- offset = nc16 * 16;
146
+ const size_t nc32 = dim / 32;
147
+ offset = nc32 * 32;
148
148
 
149
- for (size_t c = 0; c < nc16; ++c) {
150
- size_t idx = c * 16;
151
- __m512 q_vec = _mm512_loadu_ps(&query[idx]);
149
+ for (size_t c = 0; c < nc32; ++c) {
150
+ size_t idx = c * 32;
151
+ __m512 q_vec_lo = _mm512_loadu_ps(&query[idx]);
152
+ __m512 q_vec_hi = _mm512_loadu_ps(&query[idx + 16]);
152
153
 
153
154
  for (size_t j = 0; j < BATCH_SIZE; ++j) {
154
155
  const float* db_ = static_cast<const float*>(db[j]);
155
- __m512 db_vec = _mm512_loadu_ps(&db_[idx]);
156
- s[j] = _mm512_fmadd_ps(q_vec, db_vec, s[j]);
156
+ __m512 db_vec_lo = _mm512_loadu_ps(&db_[idx]);
157
+ __m512 db_vec_hi = _mm512_loadu_ps(&db_[idx + 16]);
158
+ s[j] = _mm512_fmadd_ps(q_vec_lo, db_vec_lo, s[j]);
159
+ s[j] = _mm512_fmadd_ps(q_vec_hi, db_vec_hi, s[j]);
157
160
  }
158
161
  }
159
162
  }
@@ -265,17 +268,20 @@ class InnerProductFloat_AVX2 {
265
268
  }
266
269
 
267
270
  if constexpr (HasDualSimd) {
268
- const size_t nc8 = dim / 8;
269
- offset = nc8 * 8;
271
+ const size_t nc16 = dim / 16;
272
+ offset = nc16 * 16;
270
273
 
271
- for (size_t c = 0; c < nc8; ++c) {
272
- size_t idx = c * 8;
273
- __m256 q_vec = _mm256_loadu_ps(&query[idx]);
274
+ for (size_t c = 0; c < nc16; ++c) {
275
+ size_t idx = c * 16;
276
+ __m256 q_vec_lo = _mm256_loadu_ps(&query[idx]);
277
+ __m256 q_vec_hi = _mm256_loadu_ps(&query[idx + 8]);
274
278
 
275
279
  for (size_t j = 0; j < BATCH_SIZE; ++j) {
276
280
  const float* db_ = static_cast<const float*>(db[j]);
277
- __m256 db_vec = _mm256_loadu_ps(&db_[idx]);
278
- s[j] = _mm256_fmadd_ps(q_vec, db_vec, s[j]);
281
+ __m256 db_vec_lo = _mm256_loadu_ps(&db_[idx]);
282
+ __m256 db_vec_hi = _mm256_loadu_ps(&db_[idx + 8]);
283
+ s[j] = _mm256_fmadd_ps(q_vec_lo, db_vec_lo, s[j]);
284
+ s[j] = _mm256_fmadd_ps(q_vec_hi, db_vec_hi, s[j]);
279
285
  }
280
286
  }
281
287
  }
@@ -142,17 +142,20 @@ class L2Float_AVX512 {
142
142
  }
143
143
 
144
144
  if constexpr (HasDualSimd) {
145
- const size_t nc16 = dim / 16;
146
- offset = nc16 * 16;
145
+ const size_t nc32 = dim / 32;
146
+ offset = nc32 * 32;
147
147
 
148
- for (size_t c = 0; c < nc16; ++c) {
149
- size_t idx = c * 16;
150
- __m512 q_vec = _mm512_loadu_ps(&query[idx]);
148
+ for (size_t c = 0; c < nc32; ++c) {
149
+ size_t idx = c * 32;
150
+ __m512 q_vec_lo = _mm512_loadu_ps(&query[idx]);
151
+ __m512 q_vec_hi = _mm512_loadu_ps(&query[idx + 16]);
151
152
 
152
153
  for (size_t j = 0; j < BATCH_SIZE; ++j) {
153
154
  const float* db_ = static_cast<const float*>(db[j]);
154
- __m512 diff = _mm512_sub_ps(q_vec, _mm512_loadu_ps(&db_[idx]));
155
- s[j] = _mm512_fmadd_ps(diff, diff, s[j]);
155
+ __m512 diff_lo = _mm512_sub_ps(q_vec_lo, _mm512_loadu_ps(&db_[idx]));
156
+ __m512 diff_hi = _mm512_sub_ps(q_vec_hi, _mm512_loadu_ps(&db_[idx + 16]));
157
+ s[j] = _mm512_fmadd_ps(diff_lo, diff_lo, s[j]);
158
+ s[j] = _mm512_fmadd_ps(diff_hi, diff_hi, s[j]);
156
159
  }
157
160
  }
158
161
  }
@@ -264,17 +267,20 @@ class L2Float_AVX2 {
264
267
  }
265
268
 
266
269
  if constexpr (HasDualSimd) {
267
- const size_t nc8 = dim / 8;
268
- offset = nc8 * 8;
270
+ const size_t nc16 = dim / 16;
271
+ offset = nc16 * 16;
269
272
 
270
- for (size_t c = 0; c < nc8; ++c) {
271
- size_t idx = c * 8;
272
- __m256 q_vec = _mm256_loadu_ps(&query[idx]);
273
+ for (size_t c = 0; c < nc16; ++c) {
274
+ size_t idx = c * 16;
275
+ __m256 q_vec_lo = _mm256_loadu_ps(&query[idx]);
276
+ __m256 q_vec_hi = _mm256_loadu_ps(&query[idx + 8]);
273
277
 
274
278
  for (size_t j = 0; j < BATCH_SIZE; ++j) {
275
279
  const float* db_ = static_cast<const float*>(db[j]);
276
- __m256 diff = _mm256_sub_ps(q_vec, _mm256_loadu_ps(&db_[idx]));
277
- s[j] = _mm256_fmadd_ps(diff, diff, s[j]);
280
+ __m256 diff_lo = _mm256_sub_ps(q_vec_lo, _mm256_loadu_ps(&db_[idx]));
281
+ __m256 diff_hi = _mm256_sub_ps(q_vec_hi, _mm256_loadu_ps(&db_[idx + 8]));
282
+ s[j] = _mm256_fmadd_ps(diff_lo, diff_lo, s[j]);
283
+ s[j] = _mm256_fmadd_ps(diff_hi, diff_hi, s[j]);
278
284
  }
279
285
  }
280
286
  }
@@ -169,19 +169,24 @@ class L2Int8_AVX512 {
169
169
  }
170
170
 
171
171
  if constexpr (HasDualSimd) {
172
- const size_t nc32 = dim / 32;
173
- offset = nc32 * 32;
172
+ const size_t nc64 = dim / 64;
173
+ offset = nc64 * 64;
174
174
 
175
- for (size_t c = 0; c < nc32; ++c) {
176
- size_t idx = c * 32;
177
- __m256i q_raw = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&query[idx]));
178
- __m512i q_vec = _mm512_cvtepi8_epi16(q_raw);
175
+ for (size_t c = 0; c < nc64; ++c) {
176
+ size_t idx = c * 64;
177
+ __m256i q_raw_lo = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&query[idx]));
178
+ __m256i q_raw_hi = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&query[idx + 32]));
179
+ __m512i q_vec_lo = _mm512_cvtepi8_epi16(q_raw_lo);
180
+ __m512i q_vec_hi = _mm512_cvtepi8_epi16(q_raw_hi);
179
181
 
180
182
  for (size_t j = 0; j < BATCH_SIZE; ++j) {
181
183
  const int8_t* db_ = static_cast<const int8_t*>(db[j]);
182
- __m256i r_raw = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&db_[idx]));
183
- __m512i diff = _mm512_sub_epi16(q_vec, _mm512_cvtepi8_epi16(r_raw));
184
- s[j] = _mm512_add_epi32(s[j], _mm512_madd_epi16(diff, diff));
184
+ __m256i r_raw_lo = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&db_[idx]));
185
+ __m256i r_raw_hi = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&db_[idx + 32]));
186
+ __m512i diff_lo = _mm512_sub_epi16(q_vec_lo, _mm512_cvtepi8_epi16(r_raw_lo));
187
+ __m512i diff_hi = _mm512_sub_epi16(q_vec_hi, _mm512_cvtepi8_epi16(r_raw_hi));
188
+ s[j] = _mm512_add_epi32(s[j], _mm512_madd_epi16(diff_lo, diff_lo));
189
+ s[j] = _mm512_add_epi32(s[j], _mm512_madd_epi16(diff_hi, diff_hi));
185
190
  }
186
191
  }
187
192
  }
@@ -290,19 +295,24 @@ class L2Int8_AVX2 {
290
295
  }
291
296
 
292
297
  if constexpr (HasDualSimd) {
293
- const size_t nc16 = dim / 16;
294
- offset = nc16 * 16;
298
+ const size_t nc32 = dim / 32;
299
+ offset = nc32 * 32;
295
300
 
296
- for (size_t c = 0; c < nc16; ++c) {
297
- size_t idx = c * 16;
298
- __m128i q_raw = _mm_loadu_si128(reinterpret_cast<const __m128i*>(&query[idx]));
299
- __m256i q_vec = _mm256_cvtepi8_epi16(q_raw);
301
+ for (size_t c = 0; c < nc32; ++c) {
302
+ size_t idx = c * 32;
303
+ __m128i q_raw_lo = _mm_loadu_si128(reinterpret_cast<const __m128i*>(&query[idx]));
304
+ __m128i q_raw_hi = _mm_loadu_si128(reinterpret_cast<const __m128i*>(&query[idx + 16]));
305
+ __m256i q_vec_lo = _mm256_cvtepi8_epi16(q_raw_lo);
306
+ __m256i q_vec_hi = _mm256_cvtepi8_epi16(q_raw_hi);
300
307
 
301
308
  for (size_t j = 0; j < BATCH_SIZE; ++j) {
302
309
  const int8_t* db_ = static_cast<const int8_t*>(db[j]);
303
- __m128i r_raw = _mm_loadu_si128(reinterpret_cast<const __m128i*>(&db_[idx]));
304
- __m256i diff = _mm256_sub_epi16(q_vec, _mm256_cvtepi8_epi16(r_raw));
305
- s[j] = _mm256_add_epi32(s[j], _mm256_madd_epi16(diff, diff));
310
+ __m128i r_raw_lo = _mm_loadu_si128(reinterpret_cast<const __m128i*>(&db_[idx]));
311
+ __m128i r_raw_hi = _mm_loadu_si128(reinterpret_cast<const __m128i*>(&db_[idx + 16]));
312
+ __m256i diff_lo = _mm256_sub_epi16(q_vec_lo, _mm256_cvtepi8_epi16(r_raw_lo));
313
+ __m256i diff_hi = _mm256_sub_epi16(q_vec_hi, _mm256_cvtepi8_epi16(r_raw_hi));
314
+ s[j] = _mm256_add_epi32(s[j], _mm256_madd_epi16(diff_lo, diff_lo));
315
+ s[j] = _mm256_add_epi32(s[j], _mm256_madd_epi16(diff_hi, diff_hi));
306
316
  }
307
317
  }
308
318
  }
@@ -173,19 +173,24 @@ class L2Uint8_AVX512 {
173
173
  }
174
174
 
175
175
  if constexpr (HasDualSimd) {
176
- const size_t nc32 = dim / 32;
177
- offset = nc32 * 32;
176
+ const size_t nc64 = dim / 64;
177
+ offset = nc64 * 64;
178
178
 
179
- for (size_t c = 0; c < nc32; ++c) {
180
- size_t idx = c * 32;
181
- __m256i q_raw = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&query[idx]));
182
- __m512i q_vec = _mm512_cvtepu8_epi16(q_raw);
179
+ for (size_t c = 0; c < nc64; ++c) {
180
+ size_t idx = c * 64;
181
+ __m256i q_raw_lo = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&query[idx]));
182
+ __m256i q_raw_hi = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&query[idx + 32]));
183
+ __m512i q_vec_lo = _mm512_cvtepu8_epi16(q_raw_lo);
184
+ __m512i q_vec_hi = _mm512_cvtepu8_epi16(q_raw_hi);
183
185
 
184
186
  for (size_t j = 0; j < BATCH_SIZE; ++j) {
185
187
  const unsigned char* db_ = static_cast<const unsigned char*>(db[j]);
186
- __m256i r_raw = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&db_[idx]));
187
- __m512i diff = _mm512_sub_epi16(q_vec, _mm512_cvtepu8_epi16(r_raw));
188
- s[j] = _mm512_add_epi32(s[j], _mm512_madd_epi16(diff, diff));
188
+ __m256i r_raw_lo = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&db_[idx]));
189
+ __m256i r_raw_hi = _mm256_loadu_si256(reinterpret_cast<const __m256i*>(&db_[idx + 32]));
190
+ __m512i diff_lo = _mm512_sub_epi16(q_vec_lo, _mm512_cvtepu8_epi16(r_raw_lo));
191
+ __m512i diff_hi = _mm512_sub_epi16(q_vec_hi, _mm512_cvtepu8_epi16(r_raw_hi));
192
+ s[j] = _mm512_add_epi32(s[j], _mm512_madd_epi16(diff_lo, diff_lo));
193
+ s[j] = _mm512_add_epi32(s[j], _mm512_madd_epi16(diff_hi, diff_hi));
189
194
  }
190
195
  }
191
196
  }
@@ -294,19 +299,24 @@ class L2Uint8_AVX2 {
294
299
  }
295
300
 
296
301
  if constexpr (HasDualSimd) {
297
- const size_t nc16 = dim / 16;
298
- offset = nc16 * 16;
302
+ const size_t nc32 = dim / 32;
303
+ offset = nc32 * 32;
299
304
 
300
- for (size_t c = 0; c < nc16; ++c) {
301
- size_t idx = c * 16;
302
- __m128i q_raw = _mm_loadu_si128(reinterpret_cast<const __m128i*>(&query[idx]));
303
- __m256i q_vec = _mm256_cvtepu8_epi16(q_raw);
305
+ for (size_t c = 0; c < nc32; ++c) {
306
+ size_t idx = c * 32;
307
+ __m128i q_raw_lo = _mm_loadu_si128(reinterpret_cast<const __m128i*>(&query[idx]));
308
+ __m128i q_raw_hi = _mm_loadu_si128(reinterpret_cast<const __m128i*>(&query[idx + 16]));
309
+ __m256i q_vec_lo = _mm256_cvtepu8_epi16(q_raw_lo);
310
+ __m256i q_vec_hi = _mm256_cvtepu8_epi16(q_raw_hi);
304
311
 
305
312
  for (size_t j = 0; j < BATCH_SIZE; ++j) {
306
313
  const unsigned char* db_ = static_cast<const unsigned char*>(db[j]);
307
- __m128i r_raw = _mm_loadu_si128(reinterpret_cast<const __m128i*>(&db_[idx]));
308
- __m256i diff = _mm256_sub_epi16(q_vec, _mm256_cvtepu8_epi16(r_raw));
309
- s[j] = _mm256_add_epi32(s[j], _mm256_madd_epi16(diff, diff));
314
+ __m128i r_raw_lo = _mm_loadu_si128(reinterpret_cast<const __m128i*>(&db_[idx]));
315
+ __m128i r_raw_hi = _mm_loadu_si128(reinterpret_cast<const __m128i*>(&db_[idx + 16]));
316
+ __m256i diff_lo = _mm256_sub_epi16(q_vec_lo, _mm256_cvtepu8_epi16(r_raw_lo));
317
+ __m256i diff_hi = _mm256_sub_epi16(q_vec_hi, _mm256_cvtepu8_epi16(r_raw_hi));
318
+ s[j] = _mm256_add_epi32(s[j], _mm256_madd_epi16(diff_lo, diff_lo));
319
+ s[j] = _mm256_add_epi32(s[j], _mm256_madd_epi16(diff_hi, diff_hi));
310
320
  }
311
321
  }
312
322
  }
@@ -17,17 +17,17 @@ namespace deglib::quantization::evp {
17
17
  // ============================================================================
18
18
 
19
19
  /**
20
- * Quantizes a single fp32 vector to EVP bytes.
20
+ * Quantizes a single fp32 vector directly into a pre-allocated EVP byte buffer.
21
21
  *
22
22
  * Layout: [ones (dim/8 bytes)][negative_ones (dim/8 bytes)]
23
23
  *
24
24
  * @param embedding Pointer to dim float values (row-major)
25
25
  * @param dim Dimension (must be divisible by 8)
26
26
  * @param non_zeros Number of top-K elements by absolute value
27
- * @return std::vector<std::byte> with 2 * dim/8 bytes
27
+ * @param out_bytes Pointer to buffer of at least 2 * dim/8 bytes
28
28
  * @throws std::invalid_argument if dim % 8 != 0 or non_zeros >= dim
29
29
  */
30
- inline std::vector<std::byte> quantize_single(const float* embedding, uint32_t dim, uint32_t non_zeros) {
30
+ inline void quantize_single_into(const float* embedding, uint32_t dim, uint32_t non_zeros, std::byte* out_bytes) {
31
31
  if (dim % 8 != 0) {
32
32
  throw std::invalid_argument("quantize_single: dim must be divisible by 8, got " + std::to_string(dim));
33
33
  }
@@ -36,8 +36,6 @@ inline std::vector<std::byte> quantize_single(const float* embedding, uint32_t d
36
36
  }
37
37
 
38
38
  const size_t mask_bytes = dim / 8;
39
- std::vector<std::byte> result(2 * mask_bytes);
40
-
41
39
  std::vector<std::pair<float, uint32_t>> abs_vals(dim);
42
40
  for (uint32_t i = 0; i < dim; ++i) {
43
41
  abs_vals[i] = {std::abs(embedding[i]), i};
@@ -51,8 +49,8 @@ inline std::vector<std::byte> quantize_single(const float* embedding, uint32_t d
51
49
  is_top[abs_vals[j].second] = 1;
52
50
  }
53
51
 
54
- std::byte* ones_dst = result.data();
55
- std::byte* negs_dst = result.data() + mask_bytes;
52
+ std::byte* ones_dst = out_bytes;
53
+ std::byte* negs_dst = out_bytes + mask_bytes;
56
54
 
57
55
  for (uint32_t byte_idx = 0; byte_idx < mask_bytes; ++byte_idx) {
58
56
  int byte_val = 0;
@@ -77,7 +75,23 @@ inline std::vector<std::byte> quantize_single(const float* embedding, uint32_t d
77
75
  }
78
76
  negs_dst[byte_idx] = static_cast<std::byte>(byte_val);
79
77
  }
78
+ }
80
79
 
80
+ /**
81
+ * Quantizes a single fp32 vector to EVP bytes.
82
+ *
83
+ * Layout: [ones (dim/8 bytes)][negative_ones (dim/8 bytes)]
84
+ *
85
+ * @param embedding Pointer to dim float values (row-major)
86
+ * @param dim Dimension (must be divisible by 8)
87
+ * @param non_zeros Number of top-K elements by absolute value
88
+ * @return std::vector<std::byte> with 2 * dim/8 bytes
89
+ * @throws std::invalid_argument if dim % 8 != 0 or non_zeros >= dim
90
+ */
91
+ inline std::vector<std::byte> quantize_single(const float* embedding, uint32_t dim, uint32_t non_zeros) {
92
+ const size_t mask_bytes = dim / 8;
93
+ std::vector<std::byte> result(2 * mask_bytes);
94
+ quantize_single_into(embedding, dim, non_zeros, result.data());
81
95
  return result;
82
96
  }
83
97
 
@@ -378,11 +392,69 @@ inline std::vector<std::byte> quantize_batch(const std::vector<std::vector<std::
378
392
  return result;
379
393
  }
380
394
 
395
+ /**
396
+ * Quantizes a single FP16 (uint16_t) vector directly into a pre-allocated EVP byte buffer.
397
+ *
398
+ * Layout: [ones (dim/8 bytes)][negative_ones (dim/8 bytes)]
399
+ */
400
+ inline void quantize_single_into(const uint16_t* embedding, uint32_t dim, uint32_t non_zeros, std::byte* out_bytes) {
401
+ if (dim % 8 != 0) {
402
+ throw std::invalid_argument("quantize_single: dim must be divisible by 8, got " + std::to_string(dim));
403
+ }
404
+ if (non_zeros >= dim) {
405
+ throw std::invalid_argument("quantize_single: non_zeros must be < dim");
406
+ }
407
+
408
+ const size_t mask_bytes = dim / 8;
409
+ std::vector<std::pair<uint16_t, uint32_t>> abs_vals(dim);
410
+ for (uint32_t i = 0; i < dim; ++i) {
411
+ abs_vals[i] = {static_cast<uint16_t>(embedding[i] & 0x7FFFu), i};
412
+ }
413
+ std::nth_element(abs_vals.begin(), abs_vals.begin() + (dim - non_zeros), abs_vals.end(), [](const auto& a, const auto& b) {
414
+ return a.first != b.first ? a.first < b.first : a.second < b.second;
415
+ });
416
+
417
+ std::vector<uint8_t> is_top(dim, 0);
418
+ for (uint32_t j = dim - non_zeros; j < dim; ++j) {
419
+ is_top[abs_vals[j].second] = 1;
420
+ }
421
+
422
+ std::byte* ones_dst = out_bytes;
423
+ std::byte* negs_dst = out_bytes + mask_bytes;
424
+
425
+ for (uint32_t byte_idx = 0; byte_idx < mask_bytes; ++byte_idx) {
426
+ int bv = 0;
427
+ for (int bit = 0; bit < 8; ++bit) {
428
+ const uint32_t bit_idx = byte_idx * 8 + bit;
429
+ if (bit_idx >= dim) break;
430
+ if (is_top[bit_idx] && (embedding[bit_idx] & 0x8000u) == 0 && (embedding[bit_idx] & 0x7FFFu) > 0) {
431
+ bv |= (1 << bit);
432
+ }
433
+ }
434
+ ones_dst[byte_idx] = static_cast<std::byte>(bv);
435
+ }
436
+
437
+ for (uint32_t byte_idx = 0; byte_idx < mask_bytes; ++byte_idx) {
438
+ int bv = 0;
439
+ for (int bit = 0; bit < 8; ++bit) {
440
+ const uint32_t bit_idx = byte_idx * 8 + bit;
441
+ if (bit_idx >= dim) break;
442
+ if (is_top[bit_idx] && (embedding[bit_idx] & 0x8000u) != 0 && (embedding[bit_idx] & 0x7FFFu) > 0) {
443
+ bv |= (1 << bit);
444
+ }
445
+ }
446
+ negs_dst[byte_idx] = static_cast<std::byte>(bv);
447
+ }
448
+ }
449
+
381
450
  /**
382
451
  * Quantizes a single FP16 (uint16_t) vector to EVP bytes.
383
452
  */
384
453
  inline std::vector<std::byte> quantize_single(const uint16_t* embedding, uint32_t dim, uint32_t non_zeros) {
385
- return quantize_batch(embedding, 1, dim, non_zeros, 1);
454
+ const size_t mask_bytes = dim / 8;
455
+ std::vector<std::byte> result(2 * mask_bytes);
456
+ quantize_single_into(embedding, dim, non_zeros, result.data());
457
+ return result;
386
458
  }
387
459
 
388
460
  } // namespace deglib::quantization::evp