faiss 0.6.1 → 0.6.3

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 (186) hide show
  1. checksums.yaml +4 -4
  2. data/CHANGELOG.md +8 -0
  3. data/lib/faiss/version.rb +1 -1
  4. data/vendor/faiss/faiss/AutoTune.cpp +3 -1
  5. data/vendor/faiss/faiss/Clustering.cpp +9 -1
  6. data/vendor/faiss/faiss/IVFlib.cpp +14 -3
  7. data/vendor/faiss/faiss/Index.h +2 -2
  8. data/vendor/faiss/faiss/IndexAdditiveQuantizer.cpp +9 -10
  9. data/vendor/faiss/faiss/IndexAdditiveQuantizerFastScan.cpp +2 -3
  10. data/vendor/faiss/faiss/IndexBinaryFromFloat.cpp +1 -2
  11. data/vendor/faiss/faiss/IndexBinaryHNSW.cpp +10 -12
  12. data/vendor/faiss/faiss/IndexBinaryHash.cpp +5 -9
  13. data/vendor/faiss/faiss/IndexBinaryIVF.cpp +5 -7
  14. data/vendor/faiss/faiss/IndexEDEN.cpp +273 -0
  15. data/vendor/faiss/faiss/IndexEDEN.h +57 -0
  16. data/vendor/faiss/faiss/IndexFastScan.cpp +15 -4
  17. data/vendor/faiss/faiss/IndexFlat.cpp +13 -50
  18. data/vendor/faiss/faiss/IndexHNSW.cpp +177 -148
  19. data/vendor/faiss/faiss/IndexIDMap.cpp +16 -3
  20. data/vendor/faiss/faiss/IndexIDMap.h +2 -0
  21. data/vendor/faiss/faiss/IndexIVF.cpp +19 -8
  22. data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.cpp +3 -3
  23. data/vendor/faiss/faiss/IndexIVFAdditiveQuantizerFastScan.cpp +3 -4
  24. data/vendor/faiss/faiss/IndexIVFEDEN.cpp +302 -0
  25. data/vendor/faiss/faiss/IndexIVFEDEN.h +70 -0
  26. data/vendor/faiss/faiss/IndexIVFFastScan.cpp +5 -6
  27. data/vendor/faiss/faiss/IndexIVFFlat.cpp +6 -5
  28. data/vendor/faiss/faiss/IndexIVFFlatPanorama.cpp +3 -3
  29. data/vendor/faiss/faiss/IndexIVFIndependentQuantizer.cpp +1 -1
  30. data/vendor/faiss/faiss/IndexIVFPQ.cpp +42 -25
  31. data/vendor/faiss/faiss/IndexIVFPQFastScan.cpp +0 -1
  32. data/vendor/faiss/faiss/IndexIVFPQR.cpp +2 -3
  33. data/vendor/faiss/faiss/IndexIVFRaBitQ.cpp +23 -62
  34. data/vendor/faiss/faiss/IndexIVFRaBitQFastScan.cpp +180 -76
  35. data/vendor/faiss/faiss/IndexIVFRaBitQFastScan.h +5 -4
  36. data/vendor/faiss/faiss/IndexIVFSpectralHash.cpp +8 -6
  37. data/vendor/faiss/faiss/IndexLSH.cpp +2 -3
  38. data/vendor/faiss/faiss/IndexLattice.cpp +5 -0
  39. data/vendor/faiss/faiss/IndexNNDescent.cpp +10 -3
  40. data/vendor/faiss/faiss/IndexNSG.cpp +8 -4
  41. data/vendor/faiss/faiss/IndexPQ.cpp +6 -8
  42. data/vendor/faiss/faiss/IndexPreTransform.cpp +15 -0
  43. data/vendor/faiss/faiss/IndexRaBitQ.cpp +2 -2
  44. data/vendor/faiss/faiss/IndexRaBitQFastScan.cpp +1 -2
  45. data/vendor/faiss/faiss/IndexRaBitQFastScan.h +5 -1
  46. data/vendor/faiss/faiss/IndexRefine.cpp +30 -1
  47. data/vendor/faiss/faiss/IndexReplicas.cpp +1 -2
  48. data/vendor/faiss/faiss/IndexScalarQuantizer.cpp +68 -6
  49. data/vendor/faiss/faiss/IndexScalarQuantizer.h +10 -0
  50. data/vendor/faiss/faiss/IndexShards.cpp +2 -2
  51. data/vendor/faiss/faiss/IndexShardsIVF.cpp +2 -2
  52. data/vendor/faiss/faiss/MetaIndexes.cpp +2 -4
  53. data/vendor/faiss/faiss/SuperKMeans.cpp +256 -240
  54. data/vendor/faiss/faiss/SuperKMeans.h +30 -0
  55. data/vendor/faiss/faiss/VectorTransform.cpp +33 -2
  56. data/vendor/faiss/faiss/clone_index.cpp +5 -0
  57. data/vendor/faiss/faiss/cppcontrib/SaDecodeKernels.h +1 -1
  58. data/vendor/faiss/faiss/cppcontrib/sa_decode/Level2-neon-inl.h +902 -12
  59. data/vendor/faiss/faiss/cppcontrib/sa_decode/PQ-neon-inl.h +702 -10
  60. data/vendor/faiss/faiss/factory_tools.cpp +51 -4
  61. data/vendor/faiss/faiss/gpu/GpuCloner.cpp +11 -11
  62. data/vendor/faiss/faiss/gpu/GpuIndex.h +34 -11
  63. data/vendor/faiss/faiss/gpu/GpuIndexCagra.h +47 -0
  64. data/vendor/faiss/faiss/gpu/GpuIndexIVF.h +17 -0
  65. data/vendor/faiss/faiss/gpu/GpuIndexIVFScalarQuantizer.h +16 -0
  66. data/vendor/faiss/faiss/gpu/GpuResources.h +3 -2
  67. data/vendor/faiss/faiss/gpu/StandardGpuResources.cpp +11 -12
  68. data/vendor/faiss/faiss/gpu/StandardGpuResources.h +3 -3
  69. data/vendor/faiss/faiss/gpu/perf/PerfClustering.cpp +1 -1
  70. data/vendor/faiss/faiss/gpu/perf/PerfIVFPQAdd.cpp +2 -2
  71. data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFScalarQuantizer.cpp +180 -0
  72. data/vendor/faiss/faiss/gpu_metal/MetalDistance.h +87 -0
  73. data/vendor/faiss/faiss/gpu_metal/MetalIndex.h +7 -0
  74. data/vendor/faiss/faiss/gpu_metal/MetalIndexIVFFlat.h +177 -0
  75. data/vendor/faiss/faiss/gpu_metal/MetalIndexIVFPQ.h +88 -0
  76. data/vendor/faiss/faiss/gpu_metal/MetalKernels.h +48 -3
  77. data/vendor/faiss/faiss/gpu_metal/MetalPythonBridge.h +45 -0
  78. data/vendor/faiss/faiss/gpu_metal/impl/MetalIVFFlat.h +193 -0
  79. data/vendor/faiss/faiss/gpu_metal/impl/MetalIVFPQ.h +134 -0
  80. data/vendor/faiss/faiss/impl/ClusteringInitialization.cpp +2 -2
  81. data/vendor/faiss/faiss/impl/DistanceComputer.h +34 -0
  82. data/vendor/faiss/faiss/impl/EDENQuantizer.h +119 -0
  83. data/vendor/faiss/faiss/impl/HNSW.cpp +658 -344
  84. data/vendor/faiss/faiss/impl/HNSW.h +51 -13
  85. data/vendor/faiss/faiss/impl/LocalSearchQuantizer.cpp +2 -2
  86. data/vendor/faiss/faiss/impl/NSG.cpp +18 -12
  87. data/vendor/faiss/faiss/impl/Panorama.h +20 -7
  88. data/vendor/faiss/faiss/impl/PolysemousTraining.cpp +152 -84
  89. data/vendor/faiss/faiss/impl/ProductQuantizer.cpp +59 -24
  90. data/vendor/faiss/faiss/impl/RaBitQUtils.cpp +45 -37
  91. data/vendor/faiss/faiss/impl/RaBitQUtils.h +35 -0
  92. data/vendor/faiss/faiss/impl/RaBitQuantizer.cpp +175 -68
  93. data/vendor/faiss/faiss/impl/RaBitQuantizer.h +19 -0
  94. data/vendor/faiss/faiss/impl/RaBitQuantizerMultiBit.cpp +2 -11
  95. data/vendor/faiss/faiss/impl/ResultHandler.h +26 -31
  96. data/vendor/faiss/faiss/impl/ScalarQuantizer.cpp +522 -58
  97. data/vendor/faiss/faiss/impl/ScalarQuantizer.h +70 -0
  98. data/vendor/faiss/faiss/impl/ThreadedIndex-inl.h +2 -2
  99. data/vendor/faiss/faiss/impl/VisitedTable.cpp +33 -13
  100. data/vendor/faiss/faiss/impl/VisitedTable.h +88 -33
  101. data/vendor/faiss/faiss/impl/binary_hamming/IndexBinaryIVF_impl.h +1 -1
  102. data/vendor/faiss/faiss/impl/binary_hamming/avx2.cpp +4 -4
  103. data/vendor/faiss/faiss/impl/fast_scan/dispatching.h +38 -3
  104. data/vendor/faiss/faiss/impl/hnsw/LockVector.cpp +1 -1
  105. data/vendor/faiss/faiss/impl/hnsw/MinimaxHeap.cpp +35 -43
  106. data/vendor/faiss/faiss/impl/hnsw/MinimaxHeap.h +64 -15
  107. data/vendor/faiss/faiss/impl/hnsw/avx2.cpp +86 -40
  108. data/vendor/faiss/faiss/impl/hnsw/avx512.cpp +81 -50
  109. data/vendor/faiss/faiss/impl/index_read.cpp +476 -75
  110. data/vendor/faiss/faiss/impl/index_write.cpp +56 -4
  111. data/vendor/faiss/faiss/impl/io_macros.h +25 -0
  112. data/vendor/faiss/faiss/impl/lattice_Zn.cpp +8 -9
  113. data/vendor/faiss/faiss/impl/platform_macros.h +15 -9
  114. data/vendor/faiss/faiss/impl/polysemous_training/avx512.cpp +284 -0
  115. data/vendor/faiss/faiss/impl/polysemous_training/dispatch.h +115 -0
  116. data/vendor/faiss/faiss/impl/pq_code_distance/IVFPQ_QueryTables.cpp +0 -1
  117. data/vendor/faiss/faiss/impl/pq_code_distance/PQDistanceComputer_impl.h +26 -15
  118. data/vendor/faiss/faiss/impl/pq_code_distance/avx2.cpp +6 -4
  119. data/vendor/faiss/faiss/impl/pq_code_distance/avx512.cpp +2 -0
  120. data/vendor/faiss/faiss/impl/pq_code_distance/neon.cpp +2 -0
  121. data/vendor/faiss/faiss/impl/pq_code_distance/pq_code_distance-generic.cpp +20 -0
  122. data/vendor/faiss/faiss/impl/pq_code_distance/pq_code_distance-inl.h +36 -0
  123. data/vendor/faiss/faiss/impl/pq_code_distance/pq_code_distance-sve.cpp +5 -0
  124. data/vendor/faiss/faiss/impl/pq_code_distance/pq_scan_impl.h +105 -0
  125. data/vendor/faiss/faiss/impl/pq_code_distance/rvv.cpp +2 -0
  126. data/vendor/faiss/faiss/impl/result_handler/ResultHandler.cpp +195 -0
  127. data/vendor/faiss/faiss/impl/result_handler/avx2.cpp +133 -0
  128. data/vendor/faiss/faiss/impl/result_handler/avx512.cpp +281 -0
  129. data/vendor/faiss/faiss/impl/scalar_quantizer/EDENQuantizer-avx2.cpp +72 -0
  130. data/vendor/faiss/faiss/impl/scalar_quantizer/EDENQuantizer-avx512.cpp +228 -0
  131. data/vendor/faiss/faiss/impl/scalar_quantizer/EDENQuantizer.cpp +882 -0
  132. data/vendor/faiss/faiss/impl/scalar_quantizer/distance_computers.h +6 -0
  133. data/vendor/faiss/faiss/impl/scalar_quantizer/quantizers.h +336 -26
  134. data/vendor/faiss/faiss/impl/scalar_quantizer/sq-avx2.cpp +331 -32
  135. data/vendor/faiss/faiss/impl/scalar_quantizer/sq-avx512-impl.h +553 -0
  136. data/vendor/faiss/faiss/impl/scalar_quantizer/sq-avx512-spr.cpp +558 -0
  137. data/vendor/faiss/faiss/impl/scalar_quantizer/sq-avx512.cpp +284 -45
  138. data/vendor/faiss/faiss/impl/scalar_quantizer/sq-dispatch.h +502 -3
  139. data/vendor/faiss/faiss/impl/scalar_quantizer/sq-neon.cpp +157 -32
  140. data/vendor/faiss/faiss/impl/scalar_quantizer/sq-rvv.cpp +26 -0
  141. data/vendor/faiss/faiss/impl/simd_dispatch.h +86 -8
  142. data/vendor/faiss/faiss/index_factory.cpp +37 -7
  143. data/vendor/faiss/faiss/index_io.h +16 -0
  144. data/vendor/faiss/faiss/invlists/DirectMap.cpp +5 -2
  145. data/vendor/faiss/faiss/invlists/InvertedLists.cpp +15 -15
  146. data/vendor/faiss/faiss/invlists/InvertedLists.h +2 -2
  147. data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.cpp +19 -4
  148. data/vendor/faiss/faiss/python/python_callbacks.cpp +3 -1
  149. data/vendor/faiss/faiss/svs/IndexSVSFaissUtils.h +60 -0
  150. data/vendor/faiss/faiss/svs/IndexSVSFlat.cpp +26 -1
  151. data/vendor/faiss/faiss/svs/IndexSVSFlat.h +13 -0
  152. data/vendor/faiss/faiss/svs/IndexSVSIVF.cpp +1 -1
  153. data/vendor/faiss/faiss/svs/IndexSVSIVFLeanVec.cpp +1 -1
  154. data/vendor/faiss/faiss/svs/IndexSVSVamana.cpp +150 -23
  155. data/vendor/faiss/faiss/svs/IndexSVSVamana.h +30 -7
  156. data/vendor/faiss/faiss/svs/IndexSVSVamanaLVQ.cpp +3 -2
  157. data/vendor/faiss/faiss/svs/IndexSVSVamanaLVQ.h +2 -1
  158. data/vendor/faiss/faiss/svs/IndexSVSVamanaLeanVec.cpp +65 -25
  159. data/vendor/faiss/faiss/svs/IndexSVSVamanaLeanVec.h +3 -2
  160. data/vendor/faiss/faiss/utils/approx_topk_hamming/approx_topk_hamming.h +1 -1
  161. data/vendor/faiss/faiss/utils/bf16.h +34 -0
  162. data/vendor/faiss/faiss/utils/distances.cpp +14 -2
  163. data/vendor/faiss/faiss/utils/distances_simd.cpp +4 -4
  164. data/vendor/faiss/faiss/utils/extra_distances.cpp +4 -14
  165. data/vendor/faiss/faiss/utils/extra_distances.h +1 -2
  166. data/vendor/faiss/faiss/utils/hamming.cpp +9 -9
  167. data/vendor/faiss/faiss/utils/hamming_distance/hamming_avx2.cpp +2 -1
  168. data/vendor/faiss/faiss/utils/hamming_distance/hamming_avx512_spr.cpp +15 -0
  169. data/vendor/faiss/faiss/utils/hamming_distance/hamming_computer-avx512.h +6 -30
  170. data/vendor/faiss/faiss/utils/hamming_distance/hamming_computer-avx512_spr.h +171 -0
  171. data/vendor/faiss/faiss/utils/partitioning.cpp +0 -2
  172. data/vendor/faiss/faiss/utils/quantize_lut.cpp +29 -8
  173. data/vendor/faiss/faiss/utils/rabitq_simd.h +202 -0
  174. data/vendor/faiss/faiss/utils/simd_impl/distances_avx2.cpp +0 -1
  175. data/vendor/faiss/faiss/utils/simd_impl/distances_avx512.cpp +263 -15
  176. data/vendor/faiss/faiss/utils/simd_impl/distances_rvv.cpp +160 -18
  177. data/vendor/faiss/faiss/utils/simd_impl/partitioning_simdlib256.h +14 -68
  178. data/vendor/faiss/faiss/utils/simd_impl/rabitq_avx2.cpp +245 -0
  179. data/vendor/faiss/faiss/utils/simd_impl/rabitq_avx512.cpp +273 -0
  180. data/vendor/faiss/faiss/utils/simd_impl/rabitq_avx512_spr.cpp +435 -0
  181. data/vendor/faiss/faiss/utils/simd_impl/rabitq_neon.cpp +11 -0
  182. data/vendor/faiss/faiss/utils/simd_impl/rabitq_rvv.cpp +143 -6
  183. data/vendor/faiss/faiss/utils/simd_levels.cpp +56 -2
  184. data/vendor/faiss/faiss/utils/simd_levels.h +14 -0
  185. data/vendor/faiss/faiss/utils/utils.cpp +9 -27
  186. metadata +27 -2
@@ -7,11 +7,15 @@
7
7
 
8
8
  #pragma once
9
9
 
10
+ #include <faiss/impl/RaBitQUtils.h>
10
11
  #include <faiss/impl/scalar_quantizer/codecs.h>
11
12
  #include <faiss/impl/scalar_quantizer/distance_computers.h>
12
13
  #include <faiss/impl/scalar_quantizer/quantizers.h>
13
14
  #include <faiss/impl/scalar_quantizer/scanners.h>
14
15
  #include <faiss/impl/scalar_quantizer/similarities.h>
16
+ #include <faiss/utils/distances.h>
17
+ #include <faiss/utils/rabitq_simd.h>
18
+ #include <limits>
15
19
 
16
20
  #ifndef THE_LEVEL_TO_DISPATCH
17
21
  #error "THE_LEVEL_TO_DISPATCH should be set on input to this header"
@@ -24,10 +28,324 @@ namespace scalar_quantizer {
24
28
  // Define SL as alias for THE_LEVEL_TO_DISPATCH for use in this file
25
29
  constexpr SIMDLevel SL = THE_LEVEL_TO_DISPATCH;
26
30
 
31
+ /*******************************************************************
32
+ * TurboQuant SIMD kernel: masked_sum
33
+ * Compute sum of arr[j] where bit j of the bitmask is set.
34
+ * NONE specialization is inline; AVX2/AVX512/NEON specializations
35
+ * live in sq-avx2.cpp / sq-avx512.cpp / sq-neon.cpp.
36
+ *******************************************************************/
37
+
38
+ template <SIMDLevel SL0>
39
+ float turboq_masked_sum(const float* arr, const uint8_t* bits, size_t d);
40
+
41
+ template <>
42
+ inline float turboq_masked_sum<SIMDLevel::NONE>(
43
+ const float* arr,
44
+ const uint8_t* bits,
45
+ size_t d) {
46
+ float result = 0;
47
+ for (size_t byte_idx = 0; byte_idx < (d + 7) / 8; byte_idx++) {
48
+ uint8_t b = bits[byte_idx];
49
+ size_t base = byte_idx * 8;
50
+ size_t end = std::min(base + 8, d);
51
+ for (size_t j = base; j < end; j++) {
52
+ if (b & (1 << (j - base))) {
53
+ result += arr[j];
54
+ }
55
+ }
56
+ }
57
+ return result;
58
+ }
59
+
60
+ /*******************************************************************
61
+ * Full TurboQuant DC — lives here because it needs both
62
+ * quantizers.h (QuantizerTurboQuantFull, SQTurboQFactors) and
63
+ * similarities.h (Similarity::metric_type). distance_computers.h
64
+ * can't include quantizers.h due to header ordering.
65
+ *******************************************************************/
66
+ template <int NBits, class Similarity, SIMDLevel SL2>
67
+ struct DCTurboQuantFull : ScalarQuantizer::TurboQuantRefine::DistanceComputer {
68
+ using Sim = Similarity;
69
+ QuantizerTurboQuantFull<NBits, SIMDLevel::NONE> quant;
70
+ std::vector<float> query;
71
+ std::vector<float> query_proj;
72
+ float q_norm_sq = 0;
73
+ float qjl_coeff = 0;
74
+ float total_qproj_sum = 0;
75
+
76
+ // Pre-screening state
77
+ const float* threshold_ptr = nullptr;
78
+ bool prescreen_l2 = false;
79
+ float qjl_error_coeff = 0;
80
+ mutable size_t n_total = 0;
81
+ mutable size_t n_skipped = 0;
82
+
83
+ // Integer popcount state
84
+ uint8_t qb = 0;
85
+ bool int_qjl = false;
86
+ std::vector<uint8_t> rearranged_q;
87
+ float mse_base = 0;
88
+ float mse_int_scale = 0;
89
+ float mse_popcnt_scale = 0;
90
+
91
+ // Integer QJL popcount state
92
+ std::vector<uint8_t> rearranged_qproj;
93
+ float qjl_int_scale = 0;
94
+ float qjl_popcnt_scale = 0;
95
+
96
+ // Scaled centroids for 1-bit MSE fast path (NBits==2)
97
+ float scaled_c0 = 0;
98
+ float scaled_c1 = 0;
99
+ float delta_centroid = 0;
100
+ float total_q_sum = 0;
101
+
102
+ // Multi-bit MSE decomposed coefficients (NBits==3, kMSEBits==2)
103
+ float mse_multi_base = 0;
104
+ float mse_coeff_s0 = 0;
105
+ float mse_coeff_s1 = 0;
106
+ float mse_coeff_s01 = 0;
107
+ mutable std::vector<uint8_t> scratch_and;
108
+
109
+ DCTurboQuantFull(size_t d, const std::vector<float>& trained)
110
+ : quant(d, trained) {
111
+ qjl_coeff = std::sqrt(M_PI / 2.0f) / static_cast<float>(d);
112
+ }
113
+
114
+ void configure(uint8_t qb_in, bool int_qjl_in) override {
115
+ qb = qb_in;
116
+ int_qjl = int_qjl_in;
117
+ }
118
+
119
+ void set_prescreen_threshold(const float* ptr, bool l2) override {
120
+ threshold_ptr = ptr;
121
+ prescreen_l2 = l2;
122
+ }
123
+
124
+ void clear_prescreen_threshold() override {
125
+ threshold_ptr = nullptr;
126
+ }
127
+
128
+ void set_query(const float* x) final {
129
+ q = x;
130
+ size_t d = quant.d;
131
+ query.assign(x, x + d);
132
+ q_norm_sq = fvec_norm_L2sqr(x, d);
133
+
134
+ // Project query
135
+ query_proj.resize(d);
136
+ quant.project_forward(x, query_proj.data());
137
+ float inv_sqrt_pd =
138
+ 1.0f / std::sqrt(static_cast<float>(quant.padded_d));
139
+ for (size_t j = 0; j < d; j++) {
140
+ query_proj[j] *= inv_sqrt_pd;
141
+ }
142
+
143
+ total_qproj_sum = 0;
144
+ for (size_t j = 0; j < d; j++) {
145
+ total_qproj_sum += query_proj[j];
146
+ }
147
+
148
+ // Pre-screening: worst-case L1 bound on QJL error
149
+ float qproj_l1 = 0;
150
+ for (size_t j = 0; j < d; j++) {
151
+ qproj_l1 += std::abs(query_proj[j]);
152
+ }
153
+ qjl_error_coeff = qjl_coeff * qproj_l1;
154
+
155
+ // Pre-compute for 1-bit MSE fast path
156
+ if constexpr (NBits == 2) {
157
+ float inv_sqrt_d = 1.0f / std::sqrt(static_cast<float>(d));
158
+ scaled_c0 = quant.centroids[0] * inv_sqrt_d;
159
+ scaled_c1 = quant.centroids[1] * inv_sqrt_d;
160
+ delta_centroid = scaled_c1 - scaled_c0;
161
+ total_q_sum = 0;
162
+ for (size_t j = 0; j < d; j++) {
163
+ total_q_sum += query[j];
164
+ }
165
+
166
+ // Integer popcount setup
167
+ if (qb > 0) {
168
+ size_t byte_size = (d + 7) / 8;
169
+ float q_min = *std::min_element(query.begin(), query.end());
170
+ float q_max = *std::max_element(query.begin(), query.end());
171
+ float q_range = q_max - q_min;
172
+ if (q_range < 1e-30f) {
173
+ q_range = 1e-30f;
174
+ }
175
+ float max_val = static_cast<float>((1 << qb) - 1);
176
+ float scale = max_val / q_range;
177
+ float delta_q = q_range / max_val;
178
+
179
+ rearranged_q.assign(byte_size * qb, 0);
180
+ for (size_t j = 0; j < d; j++) {
181
+ int qval = static_cast<int>(
182
+ std::round((query[j] - q_min) * scale));
183
+ qval = std::max(
184
+ 0, std::min(static_cast<int>(max_val), qval));
185
+ for (int b = 0; b < qb; b++) {
186
+ if (qval & (1 << b)) {
187
+ rearranged_q[b * byte_size + j / 8] |=
188
+ (1 << (j % 8));
189
+ }
190
+ }
191
+ }
192
+ mse_base = scaled_c0 * total_q_sum;
193
+ mse_int_scale = delta_centroid * delta_q;
194
+ mse_popcnt_scale = delta_centroid * q_min;
195
+ }
196
+ }
197
+
198
+ // Pre-compute for 2-bit MSE decomposed path (NBits==3)
199
+ if constexpr (NBits == 3) {
200
+ float inv_sqrt_d = 1.0f / std::sqrt(static_cast<float>(d));
201
+ const float* c = quant.centroids;
202
+ total_q_sum = 0;
203
+ for (size_t j = 0; j < d; j++) {
204
+ total_q_sum += query[j];
205
+ }
206
+ mse_multi_base = c[0] * inv_sqrt_d * total_q_sum;
207
+ mse_coeff_s0 = (c[1] - c[0]) * inv_sqrt_d;
208
+ mse_coeff_s1 = (c[2] - c[0]) * inv_sqrt_d;
209
+ mse_coeff_s01 = (c[3] - c[2] - c[1] + c[0]) * inv_sqrt_d;
210
+ scratch_and.resize((d + 7) / 8);
211
+ }
212
+
213
+ // Integer QJL: quantize projected query into bit-planes
214
+ if (qb > 0 && int_qjl) {
215
+ size_t byte_size = (d + 7) / 8;
216
+ float qp_min =
217
+ *std::min_element(query_proj.begin(), query_proj.end());
218
+ float qp_max =
219
+ *std::max_element(query_proj.begin(), query_proj.end());
220
+ float qp_range = qp_max - qp_min;
221
+ if (qp_range < 1e-30f) {
222
+ qp_range = 1e-30f;
223
+ }
224
+ float max_val = static_cast<float>((1 << qb) - 1);
225
+ float qp_scale = max_val / qp_range;
226
+ float delta_qp = qp_range / max_val;
227
+
228
+ rearranged_qproj.assign(byte_size * qb, 0);
229
+ for (size_t j = 0; j < d; j++) {
230
+ int qval = static_cast<int>(
231
+ std::round((query_proj[j] - qp_min) * qp_scale));
232
+ qval = std::max(0, std::min(static_cast<int>(max_val), qval));
233
+ for (int b = 0; b < qb; b++) {
234
+ if (qval & (1 << b)) {
235
+ rearranged_qproj[b * byte_size + j / 8] |=
236
+ (1 << (j % 8));
237
+ }
238
+ }
239
+ }
240
+ qjl_popcnt_scale = qp_min;
241
+ qjl_int_scale = delta_qp;
242
+ }
243
+
244
+ n_total = 0;
245
+ n_skipped = 0;
246
+ }
247
+
248
+ float query_to_code(const uint8_t* code) const final {
249
+ size_t d = quant.d;
250
+ float inv_sqrt_d = 1.0f / std::sqrt(static_cast<float>(d));
251
+ const auto* factors = reinterpret_cast<const SQTurboQFactors*>(
252
+ code + quant.mse_total_bytes + quant.qjl_plane_bytes);
253
+ float norm = factors->norm;
254
+ float gamma = factors->gamma;
255
+
256
+ // Stage 1: MSE dot product
257
+ float mse_dot = 0;
258
+ if constexpr (NBits == 2) {
259
+ if (qb > 0) {
260
+ // Integer popcount path for 1-bit MSE
261
+ size_t byte_size = (d + 7) / 8;
262
+ uint64_t and_result = rabitq::bitwise_and_dot_product<SL2>(
263
+ rearranged_q.data(), code, byte_size, qb);
264
+ uint64_t pop = rabitq::popcount<SL2>(code, byte_size);
265
+ mse_dot = mse_base +
266
+ mse_int_scale * static_cast<float>(and_result) +
267
+ mse_popcnt_scale * static_cast<float>(pop);
268
+ } else {
269
+ // Float path: masked accumulation
270
+ float pos_sum = turboq_masked_sum<SL2>(query.data(), code, d);
271
+ mse_dot = scaled_c0 * total_q_sum + delta_centroid * pos_sum;
272
+ }
273
+ } else if constexpr (NBits == 3) {
274
+ // 2-bit MSE: decompose into 3 masked sums over bit-planes.
275
+ size_t pb = quant.mse_plane_bytes;
276
+ float s0 = turboq_masked_sum<SL2>(query.data(), code, d);
277
+ float s1 = turboq_masked_sum<SL2>(query.data(), code + pb, d);
278
+ for (size_t i = 0; i < pb; i++) {
279
+ scratch_and[i] = code[i] & code[pb + i];
280
+ }
281
+ float s01 =
282
+ turboq_masked_sum<SL2>(query.data(), scratch_and.data(), d);
283
+ mse_dot = mse_multi_base + mse_coeff_s0 * s0 + mse_coeff_s1 * s1 +
284
+ mse_coeff_s01 * s01;
285
+ } else {
286
+ // kMSEBits > 2: per-dimension fallback
287
+ for (size_t j = 0; j < d; j++) {
288
+ uint8_t idx = quant.load_mse_index(code, j);
289
+ mse_dot += query[j] * quant.centroids[idx] * inv_sqrt_d;
290
+ }
291
+ }
292
+
293
+ // Pre-screening
294
+ if (threshold_ptr != nullptr) {
295
+ n_total++;
296
+ float bound = qjl_error_coeff * gamma * norm;
297
+ float mse_ip = norm * mse_dot;
298
+
299
+ if constexpr (Similarity::metric_type == METRIC_INNER_PRODUCT) {
300
+ if (mse_ip + bound <= *threshold_ptr) {
301
+ n_skipped++;
302
+ return -std::numeric_limits<float>::infinity();
303
+ }
304
+ } else {
305
+ float best_possible =
306
+ q_norm_sq + norm * norm - 2.0f * (mse_ip + bound);
307
+ if (best_possible >= *threshold_ptr) {
308
+ n_skipped++;
309
+ return std::numeric_limits<float>::infinity();
310
+ }
311
+ }
312
+ }
313
+
314
+ // Stage 2: QJL dot product
315
+ const uint8_t* qjl_code = code + quant.mse_total_bytes;
316
+ float qjl_dot;
317
+ if (qb > 0 && int_qjl) {
318
+ size_t byte_size = (d + 7) / 8;
319
+ uint64_t and_result = rabitq::bitwise_and_dot_product<SL2>(
320
+ rearranged_qproj.data(), qjl_code, byte_size, qb);
321
+ uint64_t pop = rabitq::popcount<SL2>(qjl_code, byte_size);
322
+ float pos_sum = qjl_popcnt_scale * static_cast<float>(pop) +
323
+ qjl_int_scale * static_cast<float>(and_result);
324
+ qjl_dot = qjl_coeff * gamma * (2.0f * pos_sum - total_qproj_sum);
325
+ } else {
326
+ float pos_sum =
327
+ turboq_masked_sum<SL2>(query_proj.data(), qjl_code, d);
328
+ qjl_dot = qjl_coeff * gamma * (2.0f * pos_sum - total_qproj_sum);
329
+ }
330
+
331
+ float estimated_ip = norm * (mse_dot + qjl_dot);
332
+
333
+ if constexpr (Similarity::metric_type == METRIC_INNER_PRODUCT) {
334
+ return estimated_ip;
335
+ } else {
336
+ return q_norm_sq + norm * norm - 2.0f * estimated_ip;
337
+ }
338
+ }
339
+
340
+ float symmetric_dis(idx_t, idx_t) override {
341
+ FAISS_THROW_MSG("Not implemented");
342
+ }
343
+ };
344
+
27
345
  // Returns true if dimension d is compatible with the given SIMD level
28
346
  template <SIMDLevel SL2>
29
347
  constexpr bool is_dimension_compatible(size_t d) {
30
- if constexpr (SL2 == SIMDLevel::AVX512) {
348
+ if constexpr (SL2 == SIMDLevel::AVX512 || SL2 == SIMDLevel::AVX512_SPR) {
31
349
  return d % 16 == 0;
32
350
  } else if constexpr (SL2 == SIMDLevel::AVX2 || SL2 == SIMDLevel::ARM_NEON) {
33
351
  return d % 8 == 0;
@@ -98,6 +416,42 @@ ScalarQuantizer::SQuantizer* sq_select_quantizer<THE_LEVEL_TO_DISPATCH>(
98
416
  return new QuantizerTurboQuantMSE<4, SL>(d, trained);
99
417
  case ScalarQuantizer::QT_8bit_tqmse:
100
418
  return new QuantizerTurboQuantMSE<8, SL>(d, trained);
419
+ case ScalarQuantizer::QT_1bit_eden:
420
+ return new QuantizerLloydMax<1, SL>(d, trained);
421
+ case ScalarQuantizer::QT_2bit_eden:
422
+ return new QuantizerLloydMax<2, SL>(d, trained);
423
+ case ScalarQuantizer::QT_3bit_eden:
424
+ return new QuantizerLloydMax<3, SL>(d, trained);
425
+ case ScalarQuantizer::QT_4bit_eden:
426
+ return new QuantizerLloydMax<4, SL>(d, trained);
427
+ case ScalarQuantizer::QT_5bit_eden:
428
+ if constexpr (SL != SIMDLevel::NONE) {
429
+ return nullptr;
430
+ } else {
431
+ return new QuantizerLloydMax<5, SL>(d, trained);
432
+ }
433
+ case ScalarQuantizer::QT_6bit_eden:
434
+ if constexpr (SL != SIMDLevel::NONE) {
435
+ return nullptr;
436
+ } else {
437
+ return new QuantizerLloydMax<6, SL>(d, trained);
438
+ }
439
+ case ScalarQuantizer::QT_7bit_eden:
440
+ if constexpr (SL != SIMDLevel::NONE) {
441
+ return nullptr;
442
+ } else {
443
+ return new QuantizerLloydMax<7, SL>(d, trained);
444
+ }
445
+ case ScalarQuantizer::QT_8bit_eden:
446
+ return new QuantizerLloydMax<8, SL>(d, trained);
447
+ case ScalarQuantizer::QT_2bit_tq:
448
+ return new QuantizerTurboQuantFull<2, SL>(d, trained);
449
+ case ScalarQuantizer::QT_3bit_tq:
450
+ return new QuantizerTurboQuantFull<3, SL>(d, trained);
451
+ case ScalarQuantizer::QT_4bit_tq:
452
+ return new QuantizerTurboQuantFull<4, SL>(d, trained);
453
+ case ScalarQuantizer::QT_5bit_tq:
454
+ return new QuantizerTurboQuantFull<5, SL>(d, trained);
101
455
  default:
102
456
  FAISS_THROW_MSG("unknown qtype");
103
457
  }
@@ -171,7 +525,8 @@ SQDistanceComputer* select_distance_computer_body(
171
525
  return new DCTemplate<QuantizerBF16<SL2>, Sim, SL2>(d, trained);
172
526
 
173
527
  case ScalarQuantizer::QT_8bit_direct:
174
- if constexpr (SL2 == SIMDLevel::AVX512) {
528
+ if constexpr (
529
+ SL2 == SIMDLevel::AVX512 || SL2 == SIMDLevel::AVX512_SPR) {
175
530
  if (d % 32 == 0) {
176
531
  return new DistanceComputerByte<Sim, SL2>(
177
532
  static_cast<int>(d), trained);
@@ -186,6 +541,22 @@ SQDistanceComputer* select_distance_computer_body(
186
541
  d, trained);
187
542
 
188
543
  case ScalarQuantizer::QT_8bit_direct_signed:
544
+ if constexpr (SL2 == SIMDLevel::AVX512_SPR) {
545
+ if (d % 64 == 0) {
546
+ return new DistanceComputerByteSigned<Sim, SL2>(
547
+ static_cast<int>(d), trained);
548
+ }
549
+ } else if constexpr (SL2 == SIMDLevel::AVX512) {
550
+ if (d % 32 == 0) {
551
+ return new DistanceComputerByteSigned<Sim, SL2>(
552
+ static_cast<int>(d), trained);
553
+ }
554
+ } else if constexpr (SL2 == SIMDLevel::AVX2) {
555
+ if (d % 16 == 0) {
556
+ return new DistanceComputerByteSigned<Sim, SL2>(
557
+ static_cast<int>(d), trained);
558
+ }
559
+ }
189
560
  return new DCTemplate<Quantizer8bitDirectSigned<SL2>, Sim, SL2>(
190
561
  d, trained);
191
562
  case ScalarQuantizer::QT_0bit:
@@ -206,6 +577,52 @@ SQDistanceComputer* select_distance_computer_body(
206
577
  case ScalarQuantizer::QT_8bit_tqmse:
207
578
  return new DCTemplate<QuantizerTurboQuantMSE<8, SL2>, Sim, SL2>(
208
579
  d, trained);
580
+ case ScalarQuantizer::QT_1bit_eden:
581
+ return new DCTemplate<QuantizerLloydMax<1, SL2>, Sim, SL2>(
582
+ d, trained);
583
+ case ScalarQuantizer::QT_2bit_eden:
584
+ return new DCTemplate<QuantizerLloydMax<2, SL2>, Sim, SL2>(
585
+ d, trained);
586
+ case ScalarQuantizer::QT_3bit_eden:
587
+ return new DCTemplate<QuantizerLloydMax<3, SL2>, Sim, SL2>(
588
+ d, trained);
589
+ case ScalarQuantizer::QT_4bit_eden:
590
+ return new DCTemplate<QuantizerLloydMax<4, SL2>, Sim, SL2>(
591
+ d, trained);
592
+ case ScalarQuantizer::QT_5bit_eden:
593
+ if constexpr (SL2 != SIMDLevel::NONE) {
594
+ return nullptr;
595
+ } else {
596
+ return new DCTemplate<QuantizerLloydMax<5, SL2>, Sim, SL2>(
597
+ d, trained);
598
+ }
599
+ case ScalarQuantizer::QT_6bit_eden:
600
+ if constexpr (SL2 != SIMDLevel::NONE) {
601
+ return nullptr;
602
+ } else {
603
+ return new DCTemplate<QuantizerLloydMax<6, SL2>, Sim, SL2>(
604
+ d, trained);
605
+ }
606
+ case ScalarQuantizer::QT_7bit_eden:
607
+ if constexpr (SL2 != SIMDLevel::NONE) {
608
+ return nullptr;
609
+ } else {
610
+ return new DCTemplate<QuantizerLloydMax<7, SL2>, Sim, SL2>(
611
+ d, trained);
612
+ }
613
+ case ScalarQuantizer::QT_8bit_eden:
614
+ return new DCTemplate<QuantizerLloydMax<8, SL2>, Sim, SL2>(
615
+ d, trained);
616
+ case ScalarQuantizer::QT_2bit_tq:
617
+ // FRICTION: bypasses DCTemplate entirely — custom DC
618
+ // that doesn't fit the Quantizer+Similarity decomposition
619
+ return new DCTurboQuantFull<2, Sim, SL2>(d, trained);
620
+ case ScalarQuantizer::QT_3bit_tq:
621
+ return new DCTurboQuantFull<3, Sim, SL2>(d, trained);
622
+ case ScalarQuantizer::QT_4bit_tq:
623
+ return new DCTurboQuantFull<4, Sim, SL2>(d, trained);
624
+ case ScalarQuantizer::QT_5bit_tq:
625
+ return new DCTurboQuantFull<5, Sim, SL2>(d, trained);
209
626
  default:
210
627
  FAISS_THROW_MSG("unknown qtype");
211
628
  }
@@ -320,7 +737,9 @@ InvertedListScanner* sq_select_InvertedListScanner<THE_LEVEL_TO_DISPATCH>(
320
737
  return scan.template
321
738
  operator()<DCTemplate<QuantizerBF16<SL2>, Similarity, SL2>>();
322
739
  case ScalarQuantizer::QT_8bit_direct:
323
- if constexpr (SL2 == SIMDLevel::AVX512) {
740
+ if constexpr (
741
+ SL2 == SIMDLevel::AVX512 ||
742
+ SL2 == SIMDLevel::AVX512_SPR) {
324
743
  if (d % 32 == 0) {
325
744
  return scan.template
326
745
  operator()<DistanceComputerByte<Similarity, SL2>>();
@@ -336,6 +755,22 @@ InvertedListScanner* sq_select_InvertedListScanner<THE_LEVEL_TO_DISPATCH>(
336
755
  Similarity,
337
756
  SL2>>();
338
757
  case ScalarQuantizer::QT_8bit_direct_signed:
758
+ if constexpr (SL2 == SIMDLevel::AVX512_SPR) {
759
+ if (d % 64 == 0) {
760
+ return scan.template operator()<
761
+ DistanceComputerByteSigned<Similarity, SL2>>();
762
+ }
763
+ } else if constexpr (SL2 == SIMDLevel::AVX512) {
764
+ if (d % 32 == 0) {
765
+ return scan.template operator()<
766
+ DistanceComputerByteSigned<Similarity, SL2>>();
767
+ }
768
+ } else if constexpr (SL2 == SIMDLevel::AVX2) {
769
+ if (d % 16 == 0) {
770
+ return scan.template operator()<
771
+ DistanceComputerByteSigned<Similarity, SL2>>();
772
+ }
773
+ }
339
774
  return scan.template operator()<DCTemplate<
340
775
  Quantizer8bitDirectSigned<SL2>,
341
776
  Similarity,
@@ -368,6 +803,70 @@ InvertedListScanner* sq_select_InvertedListScanner<THE_LEVEL_TO_DISPATCH>(
368
803
  QuantizerTurboQuantMSE<8, SL2>,
369
804
  Similarity,
370
805
  SL2>>();
806
+ case ScalarQuantizer::QT_1bit_eden:
807
+ return scan.template operator()<DCTemplate<
808
+ QuantizerLloydMax<1, SL2>,
809
+ Similarity,
810
+ SL2>>();
811
+ case ScalarQuantizer::QT_2bit_eden:
812
+ return scan.template operator()<DCTemplate<
813
+ QuantizerLloydMax<2, SL2>,
814
+ Similarity,
815
+ SL2>>();
816
+ case ScalarQuantizer::QT_3bit_eden:
817
+ return scan.template operator()<DCTemplate<
818
+ QuantizerLloydMax<3, SL2>,
819
+ Similarity,
820
+ SL2>>();
821
+ case ScalarQuantizer::QT_4bit_eden:
822
+ return scan.template operator()<DCTemplate<
823
+ QuantizerLloydMax<4, SL2>,
824
+ Similarity,
825
+ SL2>>();
826
+ case ScalarQuantizer::QT_5bit_eden:
827
+ if constexpr (SL2 != SIMDLevel::NONE) {
828
+ return nullptr;
829
+ } else {
830
+ return scan.template operator()<DCTemplate<
831
+ QuantizerLloydMax<5, SL2>,
832
+ Similarity,
833
+ SL2>>();
834
+ }
835
+ case ScalarQuantizer::QT_6bit_eden:
836
+ if constexpr (SL2 != SIMDLevel::NONE) {
837
+ return nullptr;
838
+ } else {
839
+ return scan.template operator()<DCTemplate<
840
+ QuantizerLloydMax<6, SL2>,
841
+ Similarity,
842
+ SL2>>();
843
+ }
844
+ case ScalarQuantizer::QT_7bit_eden:
845
+ if constexpr (SL2 != SIMDLevel::NONE) {
846
+ return nullptr;
847
+ } else {
848
+ return scan.template operator()<DCTemplate<
849
+ QuantizerLloydMax<7, SL2>,
850
+ Similarity,
851
+ SL2>>();
852
+ }
853
+ case ScalarQuantizer::QT_8bit_eden:
854
+ return scan.template operator()<DCTemplate<
855
+ QuantizerLloydMax<8, SL2>,
856
+ Similarity,
857
+ SL2>>();
858
+ case ScalarQuantizer::QT_2bit_tq:
859
+ return scan.template
860
+ operator()<DCTurboQuantFull<2, Similarity, SL2>>();
861
+ case ScalarQuantizer::QT_3bit_tq:
862
+ return scan.template
863
+ operator()<DCTurboQuantFull<3, Similarity, SL2>>();
864
+ case ScalarQuantizer::QT_4bit_tq:
865
+ return scan.template
866
+ operator()<DCTurboQuantFull<4, Similarity, SL2>>();
867
+ case ScalarQuantizer::QT_5bit_tq:
868
+ return scan.template
869
+ operator()<DCTurboQuantFull<5, Similarity, SL2>>();
371
870
  default:
372
871
  FAISS_THROW_MSG("unknown qtype");
373
872
  }