faiss 0.2.4 → 0.2.6

Sign up to get free protection for your applications and to get access to all the features.
Files changed (178) hide show
  1. checksums.yaml +4 -4
  2. data/CHANGELOG.md +10 -0
  3. data/LICENSE.txt +1 -1
  4. data/README.md +23 -21
  5. data/ext/faiss/extconf.rb +11 -0
  6. data/ext/faiss/index.cpp +17 -4
  7. data/ext/faiss/index_binary.cpp +6 -6
  8. data/ext/faiss/product_quantizer.cpp +4 -4
  9. data/lib/faiss/version.rb +1 -1
  10. data/vendor/faiss/faiss/AutoTune.cpp +13 -0
  11. data/vendor/faiss/faiss/IVFlib.cpp +101 -2
  12. data/vendor/faiss/faiss/IVFlib.h +26 -2
  13. data/vendor/faiss/faiss/Index.cpp +36 -3
  14. data/vendor/faiss/faiss/Index.h +43 -6
  15. data/vendor/faiss/faiss/Index2Layer.cpp +6 -2
  16. data/vendor/faiss/faiss/Index2Layer.h +6 -1
  17. data/vendor/faiss/faiss/IndexAdditiveQuantizer.cpp +219 -16
  18. data/vendor/faiss/faiss/IndexAdditiveQuantizer.h +63 -5
  19. data/vendor/faiss/faiss/IndexAdditiveQuantizerFastScan.cpp +299 -0
  20. data/vendor/faiss/faiss/IndexAdditiveQuantizerFastScan.h +199 -0
  21. data/vendor/faiss/faiss/IndexBinary.cpp +20 -4
  22. data/vendor/faiss/faiss/IndexBinary.h +18 -3
  23. data/vendor/faiss/faiss/IndexBinaryFlat.cpp +9 -2
  24. data/vendor/faiss/faiss/IndexBinaryFlat.h +4 -2
  25. data/vendor/faiss/faiss/IndexBinaryFromFloat.cpp +4 -1
  26. data/vendor/faiss/faiss/IndexBinaryFromFloat.h +2 -1
  27. data/vendor/faiss/faiss/IndexBinaryHNSW.cpp +5 -1
  28. data/vendor/faiss/faiss/IndexBinaryHNSW.h +2 -1
  29. data/vendor/faiss/faiss/IndexBinaryHash.cpp +17 -4
  30. data/vendor/faiss/faiss/IndexBinaryHash.h +8 -4
  31. data/vendor/faiss/faiss/IndexBinaryIVF.cpp +28 -13
  32. data/vendor/faiss/faiss/IndexBinaryIVF.h +10 -7
  33. data/vendor/faiss/faiss/IndexFastScan.cpp +626 -0
  34. data/vendor/faiss/faiss/IndexFastScan.h +145 -0
  35. data/vendor/faiss/faiss/IndexFlat.cpp +34 -21
  36. data/vendor/faiss/faiss/IndexFlat.h +7 -4
  37. data/vendor/faiss/faiss/IndexFlatCodes.cpp +35 -1
  38. data/vendor/faiss/faiss/IndexFlatCodes.h +12 -0
  39. data/vendor/faiss/faiss/IndexHNSW.cpp +66 -138
  40. data/vendor/faiss/faiss/IndexHNSW.h +4 -2
  41. data/vendor/faiss/faiss/IndexIDMap.cpp +247 -0
  42. data/vendor/faiss/faiss/IndexIDMap.h +107 -0
  43. data/vendor/faiss/faiss/IndexIVF.cpp +121 -33
  44. data/vendor/faiss/faiss/IndexIVF.h +35 -16
  45. data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.cpp +84 -7
  46. data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.h +63 -1
  47. data/vendor/faiss/faiss/IndexIVFAdditiveQuantizerFastScan.cpp +590 -0
  48. data/vendor/faiss/faiss/IndexIVFAdditiveQuantizerFastScan.h +171 -0
  49. data/vendor/faiss/faiss/IndexIVFFastScan.cpp +1290 -0
  50. data/vendor/faiss/faiss/IndexIVFFastScan.h +213 -0
  51. data/vendor/faiss/faiss/IndexIVFFlat.cpp +37 -17
  52. data/vendor/faiss/faiss/IndexIVFFlat.h +4 -2
  53. data/vendor/faiss/faiss/IndexIVFPQ.cpp +234 -50
  54. data/vendor/faiss/faiss/IndexIVFPQ.h +5 -1
  55. data/vendor/faiss/faiss/IndexIVFPQFastScan.cpp +23 -852
  56. data/vendor/faiss/faiss/IndexIVFPQFastScan.h +7 -112
  57. data/vendor/faiss/faiss/IndexIVFPQR.cpp +3 -3
  58. data/vendor/faiss/faiss/IndexIVFPQR.h +1 -1
  59. data/vendor/faiss/faiss/IndexIVFSpectralHash.cpp +3 -1
  60. data/vendor/faiss/faiss/IndexIVFSpectralHash.h +2 -1
  61. data/vendor/faiss/faiss/IndexLSH.cpp +4 -2
  62. data/vendor/faiss/faiss/IndexLSH.h +2 -1
  63. data/vendor/faiss/faiss/IndexLattice.cpp +7 -1
  64. data/vendor/faiss/faiss/IndexLattice.h +3 -1
  65. data/vendor/faiss/faiss/IndexNNDescent.cpp +4 -3
  66. data/vendor/faiss/faiss/IndexNNDescent.h +2 -1
  67. data/vendor/faiss/faiss/IndexNSG.cpp +37 -3
  68. data/vendor/faiss/faiss/IndexNSG.h +25 -1
  69. data/vendor/faiss/faiss/IndexPQ.cpp +106 -69
  70. data/vendor/faiss/faiss/IndexPQ.h +19 -5
  71. data/vendor/faiss/faiss/IndexPQFastScan.cpp +15 -450
  72. data/vendor/faiss/faiss/IndexPQFastScan.h +15 -78
  73. data/vendor/faiss/faiss/IndexPreTransform.cpp +47 -8
  74. data/vendor/faiss/faiss/IndexPreTransform.h +15 -3
  75. data/vendor/faiss/faiss/IndexRefine.cpp +8 -4
  76. data/vendor/faiss/faiss/IndexRefine.h +4 -2
  77. data/vendor/faiss/faiss/IndexReplicas.cpp +4 -2
  78. data/vendor/faiss/faiss/IndexReplicas.h +2 -1
  79. data/vendor/faiss/faiss/IndexRowwiseMinMax.cpp +438 -0
  80. data/vendor/faiss/faiss/IndexRowwiseMinMax.h +92 -0
  81. data/vendor/faiss/faiss/IndexScalarQuantizer.cpp +26 -15
  82. data/vendor/faiss/faiss/IndexScalarQuantizer.h +6 -7
  83. data/vendor/faiss/faiss/IndexShards.cpp +4 -1
  84. data/vendor/faiss/faiss/IndexShards.h +2 -1
  85. data/vendor/faiss/faiss/MetaIndexes.cpp +5 -178
  86. data/vendor/faiss/faiss/MetaIndexes.h +3 -81
  87. data/vendor/faiss/faiss/VectorTransform.cpp +43 -0
  88. data/vendor/faiss/faiss/VectorTransform.h +22 -4
  89. data/vendor/faiss/faiss/clone_index.cpp +23 -1
  90. data/vendor/faiss/faiss/clone_index.h +3 -0
  91. data/vendor/faiss/faiss/cppcontrib/SaDecodeKernels.h +300 -0
  92. data/vendor/faiss/faiss/cppcontrib/detail/CoarseBitType.h +24 -0
  93. data/vendor/faiss/faiss/cppcontrib/detail/UintReader.h +195 -0
  94. data/vendor/faiss/faiss/cppcontrib/sa_decode/Level2-avx2-inl.h +2058 -0
  95. data/vendor/faiss/faiss/cppcontrib/sa_decode/Level2-inl.h +408 -0
  96. data/vendor/faiss/faiss/cppcontrib/sa_decode/Level2-neon-inl.h +2147 -0
  97. data/vendor/faiss/faiss/cppcontrib/sa_decode/MinMax-inl.h +460 -0
  98. data/vendor/faiss/faiss/cppcontrib/sa_decode/MinMaxFP16-inl.h +465 -0
  99. data/vendor/faiss/faiss/cppcontrib/sa_decode/PQ-avx2-inl.h +1618 -0
  100. data/vendor/faiss/faiss/cppcontrib/sa_decode/PQ-inl.h +251 -0
  101. data/vendor/faiss/faiss/cppcontrib/sa_decode/PQ-neon-inl.h +1452 -0
  102. data/vendor/faiss/faiss/gpu/GpuAutoTune.cpp +1 -0
  103. data/vendor/faiss/faiss/gpu/GpuCloner.cpp +0 -4
  104. data/vendor/faiss/faiss/gpu/GpuIndex.h +28 -4
  105. data/vendor/faiss/faiss/gpu/GpuIndexBinaryFlat.h +2 -1
  106. data/vendor/faiss/faiss/gpu/GpuIndexFlat.h +10 -8
  107. data/vendor/faiss/faiss/gpu/GpuIndexIVF.h +75 -14
  108. data/vendor/faiss/faiss/gpu/GpuIndexIVFFlat.h +19 -32
  109. data/vendor/faiss/faiss/gpu/GpuIndexIVFPQ.h +22 -31
  110. data/vendor/faiss/faiss/gpu/GpuIndexIVFScalarQuantizer.h +22 -28
  111. data/vendor/faiss/faiss/gpu/GpuResources.cpp +14 -0
  112. data/vendor/faiss/faiss/gpu/GpuResources.h +16 -3
  113. data/vendor/faiss/faiss/gpu/StandardGpuResources.cpp +3 -3
  114. data/vendor/faiss/faiss/gpu/impl/IndexUtils.h +32 -0
  115. data/vendor/faiss/faiss/gpu/test/TestGpuIndexBinaryFlat.cpp +1 -0
  116. data/vendor/faiss/faiss/gpu/test/TestGpuIndexFlat.cpp +311 -75
  117. data/vendor/faiss/faiss/gpu/test/TestUtils.cpp +10 -0
  118. data/vendor/faiss/faiss/gpu/test/TestUtils.h +3 -0
  119. data/vendor/faiss/faiss/gpu/test/demo_ivfpq_indexing_gpu.cpp +2 -2
  120. data/vendor/faiss/faiss/gpu/utils/DeviceUtils.h +5 -4
  121. data/vendor/faiss/faiss/impl/AdditiveQuantizer.cpp +116 -47
  122. data/vendor/faiss/faiss/impl/AdditiveQuantizer.h +44 -13
  123. data/vendor/faiss/faiss/impl/AuxIndexStructures.cpp +0 -54
  124. data/vendor/faiss/faiss/impl/AuxIndexStructures.h +0 -76
  125. data/vendor/faiss/faiss/impl/DistanceComputer.h +64 -0
  126. data/vendor/faiss/faiss/impl/HNSW.cpp +123 -27
  127. data/vendor/faiss/faiss/impl/HNSW.h +19 -16
  128. data/vendor/faiss/faiss/impl/IDSelector.cpp +125 -0
  129. data/vendor/faiss/faiss/impl/IDSelector.h +135 -0
  130. data/vendor/faiss/faiss/impl/LocalSearchQuantizer.cpp +6 -28
  131. data/vendor/faiss/faiss/impl/LocalSearchQuantizer.h +6 -1
  132. data/vendor/faiss/faiss/impl/LookupTableScaler.h +77 -0
  133. data/vendor/faiss/faiss/impl/NNDescent.cpp +1 -0
  134. data/vendor/faiss/faiss/impl/NSG.cpp +1 -1
  135. data/vendor/faiss/faiss/impl/ProductAdditiveQuantizer.cpp +383 -0
  136. data/vendor/faiss/faiss/impl/ProductAdditiveQuantizer.h +154 -0
  137. data/vendor/faiss/faiss/impl/ProductQuantizer.cpp +225 -145
  138. data/vendor/faiss/faiss/impl/ProductQuantizer.h +29 -10
  139. data/vendor/faiss/faiss/impl/Quantizer.h +43 -0
  140. data/vendor/faiss/faiss/impl/ResidualQuantizer.cpp +192 -36
  141. data/vendor/faiss/faiss/impl/ResidualQuantizer.h +40 -20
  142. data/vendor/faiss/faiss/impl/ResultHandler.h +96 -0
  143. data/vendor/faiss/faiss/impl/ScalarQuantizer.cpp +97 -173
  144. data/vendor/faiss/faiss/impl/ScalarQuantizer.h +18 -18
  145. data/vendor/faiss/faiss/impl/index_read.cpp +240 -9
  146. data/vendor/faiss/faiss/impl/index_write.cpp +237 -5
  147. data/vendor/faiss/faiss/impl/kmeans1d.cpp +6 -4
  148. data/vendor/faiss/faiss/impl/pq4_fast_scan.cpp +56 -16
  149. data/vendor/faiss/faiss/impl/pq4_fast_scan.h +25 -8
  150. data/vendor/faiss/faiss/impl/pq4_fast_scan_search_1.cpp +66 -25
  151. data/vendor/faiss/faiss/impl/pq4_fast_scan_search_qbs.cpp +75 -27
  152. data/vendor/faiss/faiss/index_factory.cpp +196 -7
  153. data/vendor/faiss/faiss/index_io.h +5 -0
  154. data/vendor/faiss/faiss/invlists/DirectMap.cpp +1 -0
  155. data/vendor/faiss/faiss/invlists/InvertedLists.cpp +4 -1
  156. data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.cpp +2 -1
  157. data/vendor/faiss/faiss/python/python_callbacks.cpp +27 -0
  158. data/vendor/faiss/faiss/python/python_callbacks.h +15 -0
  159. data/vendor/faiss/faiss/utils/Heap.h +31 -15
  160. data/vendor/faiss/faiss/utils/distances.cpp +380 -56
  161. data/vendor/faiss/faiss/utils/distances.h +113 -15
  162. data/vendor/faiss/faiss/utils/distances_simd.cpp +726 -6
  163. data/vendor/faiss/faiss/utils/extra_distances.cpp +12 -7
  164. data/vendor/faiss/faiss/utils/extra_distances.h +3 -1
  165. data/vendor/faiss/faiss/utils/fp16-fp16c.h +21 -0
  166. data/vendor/faiss/faiss/utils/fp16-inl.h +101 -0
  167. data/vendor/faiss/faiss/utils/fp16.h +11 -0
  168. data/vendor/faiss/faiss/utils/hamming-inl.h +54 -0
  169. data/vendor/faiss/faiss/utils/hamming.cpp +0 -48
  170. data/vendor/faiss/faiss/utils/ordered_key_value.h +10 -0
  171. data/vendor/faiss/faiss/utils/quantize_lut.cpp +62 -0
  172. data/vendor/faiss/faiss/utils/quantize_lut.h +20 -0
  173. data/vendor/faiss/faiss/utils/random.cpp +53 -0
  174. data/vendor/faiss/faiss/utils/random.h +5 -0
  175. data/vendor/faiss/faiss/utils/simdlib_avx2.h +4 -0
  176. data/vendor/faiss/faiss/utils/simdlib_emulated.h +6 -1
  177. data/vendor/faiss/faiss/utils/simdlib_neon.h +7 -2
  178. metadata +37 -3
@@ -0,0 +1,299 @@
1
+ /**
2
+ * Copyright (c) Facebook, Inc. and its affiliates.
3
+ *
4
+ * This source code is licensed under the MIT license found in the
5
+ * LICENSE file in the root directory of this source tree.
6
+ */
7
+
8
+ #include <faiss/IndexAdditiveQuantizerFastScan.h>
9
+
10
+ #include <limits.h>
11
+ #include <cassert>
12
+ #include <memory>
13
+
14
+ #include <omp.h>
15
+
16
+ #include <faiss/impl/FaissAssert.h>
17
+ #include <faiss/impl/LocalSearchQuantizer.h>
18
+ #include <faiss/impl/LookupTableScaler.h>
19
+ #include <faiss/impl/ResidualQuantizer.h>
20
+ #include <faiss/impl/pq4_fast_scan.h>
21
+ #include <faiss/utils/quantize_lut.h>
22
+ #include <faiss/utils/utils.h>
23
+
24
+ namespace faiss {
25
+
26
+ inline size_t roundup(size_t a, size_t b) {
27
+ return (a + b - 1) / b * b;
28
+ }
29
+
30
+ IndexAdditiveQuantizerFastScan::IndexAdditiveQuantizerFastScan(
31
+ AdditiveQuantizer* aq,
32
+ MetricType metric,
33
+ int bbs) {
34
+ init(aq, metric, bbs);
35
+ }
36
+
37
+ void IndexAdditiveQuantizerFastScan::init(
38
+ AdditiveQuantizer* aq,
39
+ MetricType metric,
40
+ int bbs) {
41
+ FAISS_THROW_IF_NOT(aq != nullptr);
42
+ FAISS_THROW_IF_NOT(!aq->nbits.empty());
43
+ FAISS_THROW_IF_NOT(aq->nbits[0] == 4);
44
+ if (metric == METRIC_INNER_PRODUCT) {
45
+ FAISS_THROW_IF_NOT_MSG(
46
+ aq->search_type == AdditiveQuantizer::ST_LUT_nonorm,
47
+ "Search type must be ST_LUT_nonorm for IP metric");
48
+ } else {
49
+ FAISS_THROW_IF_NOT_MSG(
50
+ aq->search_type == AdditiveQuantizer::ST_norm_lsq2x4 ||
51
+ aq->search_type == AdditiveQuantizer::ST_norm_rq2x4,
52
+ "Search type must be lsq2x4 or rq2x4 for L2 metric");
53
+ }
54
+
55
+ this->aq = aq;
56
+ if (metric == METRIC_L2) {
57
+ M = aq->M + 2; // 2x4 bits AQ
58
+ } else {
59
+ M = aq->M;
60
+ }
61
+ init_fastscan(aq->d, M, 4, metric, bbs);
62
+
63
+ max_train_points = 1024 * ksub * M;
64
+ }
65
+
66
+ IndexAdditiveQuantizerFastScan::IndexAdditiveQuantizerFastScan()
67
+ : IndexFastScan() {
68
+ is_trained = false;
69
+ aq = nullptr;
70
+ }
71
+
72
+ IndexAdditiveQuantizerFastScan::IndexAdditiveQuantizerFastScan(
73
+ const IndexAdditiveQuantizer& orig,
74
+ int bbs) {
75
+ init(orig.aq, orig.metric_type, bbs);
76
+
77
+ ntotal = orig.ntotal;
78
+ is_trained = orig.is_trained;
79
+ orig_codes = orig.codes.data();
80
+
81
+ ntotal2 = roundup(ntotal, bbs);
82
+ codes.resize(ntotal2 * M2 / 2);
83
+ pq4_pack_codes(orig_codes, ntotal, M, ntotal2, bbs, M2, codes.get());
84
+ }
85
+
86
+ IndexAdditiveQuantizerFastScan::~IndexAdditiveQuantizerFastScan() {}
87
+
88
+ void IndexAdditiveQuantizerFastScan::train(idx_t n, const float* x_in) {
89
+ if (is_trained) {
90
+ return;
91
+ }
92
+
93
+ const int seed = 0x12345;
94
+ size_t nt = n;
95
+ const float* x = fvecs_maybe_subsample(
96
+ d, &nt, max_train_points, x_in, verbose, seed);
97
+ n = nt;
98
+ if (verbose) {
99
+ printf("training additive quantizer on %zd vectors\n", nt);
100
+ }
101
+
102
+ aq->verbose = verbose;
103
+ aq->train(n, x);
104
+ if (metric_type == METRIC_L2) {
105
+ estimate_norm_scale(n, x);
106
+ }
107
+
108
+ is_trained = true;
109
+ }
110
+
111
+ void IndexAdditiveQuantizerFastScan::estimate_norm_scale(
112
+ idx_t n,
113
+ const float* x_in) {
114
+ FAISS_THROW_IF_NOT(metric_type == METRIC_L2);
115
+
116
+ constexpr int seed = 0x980903;
117
+ constexpr size_t max_points_estimated = 65536;
118
+ size_t ns = n;
119
+ const float* x = fvecs_maybe_subsample(
120
+ d, &ns, max_points_estimated, x_in, verbose, seed);
121
+ n = ns;
122
+ std::unique_ptr<float[]> del_x;
123
+ if (x != x_in) {
124
+ del_x.reset((float*)x);
125
+ }
126
+
127
+ std::vector<float> dis_tables(n * M * ksub);
128
+ compute_float_LUT(dis_tables.data(), n, x);
129
+
130
+ // here we compute the mean of scales for each query
131
+ // TODO: try max of scales
132
+ double scale = 0;
133
+
134
+ #pragma omp parallel for reduction(+ : scale)
135
+ for (idx_t i = 0; i < n; i++) {
136
+ const float* lut = dis_tables.data() + i * M * ksub;
137
+ scale += quantize_lut::aq_estimate_norm_scale(M, ksub, 2, lut);
138
+ }
139
+ scale /= n;
140
+ norm_scale = (int)std::roundf(std::max(scale, 1.0));
141
+
142
+ if (verbose) {
143
+ printf("estimated norm scale: %lf\n", scale);
144
+ printf("rounded norm scale: %d\n", norm_scale);
145
+ }
146
+ }
147
+
148
+ void IndexAdditiveQuantizerFastScan::compute_codes(
149
+ uint8_t* tmp_codes,
150
+ idx_t n,
151
+ const float* x) const {
152
+ aq->compute_codes(x, tmp_codes, n);
153
+ }
154
+
155
+ void IndexAdditiveQuantizerFastScan::compute_float_LUT(
156
+ float* lut,
157
+ idx_t n,
158
+ const float* x) const {
159
+ if (metric_type == METRIC_INNER_PRODUCT) {
160
+ aq->compute_LUT(n, x, lut, 1.0f);
161
+ } else {
162
+ // compute inner product look-up tables
163
+ const size_t ip_dim12 = aq->M * ksub;
164
+ const size_t norm_dim12 = 2 * ksub;
165
+ std::vector<float> ip_lut(n * ip_dim12);
166
+ aq->compute_LUT(n, x, ip_lut.data(), -2.0f);
167
+
168
+ // copy and rescale norm look-up tables
169
+ auto norm_tabs = aq->norm_tabs;
170
+ if (rescale_norm && norm_scale > 1 && metric_type == METRIC_L2) {
171
+ for (size_t i = 0; i < norm_tabs.size(); i++) {
172
+ norm_tabs[i] /= norm_scale;
173
+ }
174
+ }
175
+ const float* norm_lut = norm_tabs.data();
176
+ FAISS_THROW_IF_NOT(norm_tabs.size() == norm_dim12);
177
+
178
+ // combine them
179
+ for (idx_t i = 0; i < n; i++) {
180
+ memcpy(lut, ip_lut.data() + i * ip_dim12, ip_dim12 * sizeof(*lut));
181
+ lut += ip_dim12;
182
+ memcpy(lut, norm_lut, norm_dim12 * sizeof(*lut));
183
+ lut += norm_dim12;
184
+ }
185
+ }
186
+ }
187
+
188
+ void IndexAdditiveQuantizerFastScan::search(
189
+ idx_t n,
190
+ const float* x,
191
+ idx_t k,
192
+ float* distances,
193
+ idx_t* labels,
194
+ const SearchParameters* params) const {
195
+ FAISS_THROW_IF_NOT_MSG(
196
+ !params, "search params not supported for this index");
197
+ FAISS_THROW_IF_NOT(k > 0);
198
+ bool rescale = (rescale_norm && norm_scale > 1 && metric_type == METRIC_L2);
199
+ if (!rescale) {
200
+ IndexFastScan::search(n, x, k, distances, labels);
201
+ return;
202
+ }
203
+
204
+ NormTableScaler scaler(norm_scale);
205
+ if (metric_type == METRIC_L2) {
206
+ search_dispatch_implem<true>(n, x, k, distances, labels, scaler);
207
+ } else {
208
+ search_dispatch_implem<false>(n, x, k, distances, labels, scaler);
209
+ }
210
+ }
211
+
212
+ void IndexAdditiveQuantizerFastScan::sa_decode(
213
+ idx_t n,
214
+ const uint8_t* bytes,
215
+ float* x) const {
216
+ aq->decode(bytes, x, n);
217
+ }
218
+
219
+ /**************************************************************************************
220
+ * IndexResidualQuantizerFastScan
221
+ **************************************************************************************/
222
+
223
+ IndexResidualQuantizerFastScan::IndexResidualQuantizerFastScan(
224
+ int d, ///< dimensionality of the input vectors
225
+ size_t M, ///< number of subquantizers
226
+ size_t nbits, ///< number of bit per subvector index
227
+ MetricType metric,
228
+ Search_type_t search_type,
229
+ int bbs)
230
+ : rq(d, M, nbits, search_type) {
231
+ init(&rq, metric, bbs);
232
+ }
233
+
234
+ IndexResidualQuantizerFastScan::IndexResidualQuantizerFastScan() {
235
+ aq = &rq;
236
+ }
237
+
238
+ /**************************************************************************************
239
+ * IndexLocalSearchQuantizerFastScan
240
+ **************************************************************************************/
241
+
242
+ IndexLocalSearchQuantizerFastScan::IndexLocalSearchQuantizerFastScan(
243
+ int d,
244
+ size_t M, ///< number of subquantizers
245
+ size_t nbits, ///< number of bit per subvector index
246
+ MetricType metric,
247
+ Search_type_t search_type,
248
+ int bbs)
249
+ : lsq(d, M, nbits, search_type) {
250
+ init(&lsq, metric, bbs);
251
+ }
252
+
253
+ IndexLocalSearchQuantizerFastScan::IndexLocalSearchQuantizerFastScan() {
254
+ aq = &lsq;
255
+ }
256
+
257
+ /**************************************************************************************
258
+ * IndexProductResidualQuantizerFastScan
259
+ **************************************************************************************/
260
+
261
+ IndexProductResidualQuantizerFastScan::IndexProductResidualQuantizerFastScan(
262
+ int d, ///< dimensionality of the input vectors
263
+ size_t nsplits, ///< number of residual quantizers
264
+ size_t Msub, ///< number of subquantizers per RQ
265
+ size_t nbits, ///< number of bit per subvector index
266
+ MetricType metric,
267
+ Search_type_t search_type,
268
+ int bbs)
269
+ : prq(d, nsplits, Msub, nbits, search_type) {
270
+ init(&prq, metric, bbs);
271
+ }
272
+
273
+ IndexProductResidualQuantizerFastScan::IndexProductResidualQuantizerFastScan() {
274
+ aq = &prq;
275
+ }
276
+
277
+ /**************************************************************************************
278
+ * IndexProductLocalSearchQuantizerFastScan
279
+ **************************************************************************************/
280
+
281
+ IndexProductLocalSearchQuantizerFastScan::
282
+ IndexProductLocalSearchQuantizerFastScan(
283
+ int d, ///< dimensionality of the input vectors
284
+ size_t nsplits, ///< number of local search quantizers
285
+ size_t Msub, ///< number of subquantizers per LSQ
286
+ size_t nbits, ///< number of bit per subvector index
287
+ MetricType metric,
288
+ Search_type_t search_type,
289
+ int bbs)
290
+ : plsq(d, nsplits, Msub, nbits, search_type) {
291
+ init(&plsq, metric, bbs);
292
+ }
293
+
294
+ IndexProductLocalSearchQuantizerFastScan::
295
+ IndexProductLocalSearchQuantizerFastScan() {
296
+ aq = &plsq;
297
+ }
298
+
299
+ } // namespace faiss
@@ -0,0 +1,199 @@
1
+ /**
2
+ * Copyright (c) Facebook, Inc. and its affiliates.
3
+ *
4
+ * This source code is licensed under the MIT license found in the
5
+ * LICENSE file in the root directory of this source tree.
6
+ */
7
+
8
+ #pragma once
9
+
10
+ #include <faiss/IndexAdditiveQuantizer.h>
11
+ #include <faiss/IndexFastScan.h>
12
+ #include <faiss/impl/AdditiveQuantizer.h>
13
+ #include <faiss/impl/ProductAdditiveQuantizer.h>
14
+ #include <faiss/utils/AlignedTable.h>
15
+
16
+ namespace faiss {
17
+
18
+ /** Fast scan version of IndexAQ. Works for 4-bit AQ for now.
19
+ *
20
+ * The codes are not stored sequentially but grouped in blocks of size bbs.
21
+ * This makes it possible to compute distances quickly with SIMD instructions.
22
+ *
23
+ * Implementations:
24
+ * 12: blocked loop with internal loop on Q with qbs
25
+ * 13: same with reservoir accumulator to store results
26
+ * 14: no qbs with heap accumulator
27
+ * 15: no qbs with reservoir accumulator
28
+ */
29
+
30
+ struct IndexAdditiveQuantizerFastScan : IndexFastScan {
31
+ AdditiveQuantizer* aq;
32
+ using Search_type_t = AdditiveQuantizer::Search_type_t;
33
+
34
+ bool rescale_norm = true;
35
+ int norm_scale = 1;
36
+
37
+ // max number of training vectors
38
+ size_t max_train_points = 0;
39
+
40
+ explicit IndexAdditiveQuantizerFastScan(
41
+ AdditiveQuantizer* aq,
42
+ MetricType metric = METRIC_L2,
43
+ int bbs = 32);
44
+
45
+ void init(
46
+ AdditiveQuantizer* aq,
47
+ MetricType metric = METRIC_L2,
48
+ int bbs = 32);
49
+
50
+ IndexAdditiveQuantizerFastScan();
51
+
52
+ ~IndexAdditiveQuantizerFastScan() override;
53
+
54
+ /// build from an existing IndexAQ
55
+ explicit IndexAdditiveQuantizerFastScan(
56
+ const IndexAdditiveQuantizer& orig,
57
+ int bbs = 32);
58
+
59
+ void train(idx_t n, const float* x) override;
60
+
61
+ void estimate_norm_scale(idx_t n, const float* x);
62
+
63
+ void compute_codes(uint8_t* codes, idx_t n, const float* x) const override;
64
+
65
+ void compute_float_LUT(float* lut, idx_t n, const float* x) const override;
66
+
67
+ void search(
68
+ idx_t n,
69
+ const float* x,
70
+ idx_t k,
71
+ float* distances,
72
+ idx_t* labels,
73
+ const SearchParameters* params = nullptr) const override;
74
+
75
+ /** Decode a set of vectors.
76
+ *
77
+ * NOTE: The codes in the IndexAdditiveQuantizerFastScan object are non-
78
+ * contiguous. But this method requires a contiguous representation.
79
+ *
80
+ * @param n number of vectors
81
+ * @param bytes input encoded vectors, size n * code_size
82
+ * @param x output vectors, size n * d
83
+ */
84
+ void sa_decode(idx_t n, const uint8_t* bytes, float* x) const override;
85
+ };
86
+
87
+ /** Index based on a residual quantizer. Stored vectors are
88
+ * approximated by residual quantization codes.
89
+ * Can also be used as a codec
90
+ */
91
+ struct IndexResidualQuantizerFastScan : IndexAdditiveQuantizerFastScan {
92
+ /// The residual quantizer used to encode the vectors
93
+ ResidualQuantizer rq;
94
+
95
+ /** Constructor.
96
+ *
97
+ * @param d dimensionality of the input vectors
98
+ * @param M number of subquantizers
99
+ * @param nbits number of bit per subvector index
100
+ * @param metric metric type
101
+ * @param search_type AQ search type
102
+ */
103
+ IndexResidualQuantizerFastScan(
104
+ int d, ///< dimensionality of the input vectors
105
+ size_t M, ///< number of subquantizers
106
+ size_t nbits, ///< number of bit per subvector index
107
+ MetricType metric = METRIC_L2,
108
+ Search_type_t search_type = AdditiveQuantizer::ST_norm_rq2x4,
109
+ int bbs = 32);
110
+
111
+ IndexResidualQuantizerFastScan();
112
+ };
113
+
114
+ /** Index based on a local search quantizer. Stored vectors are
115
+ * approximated by local search quantization codes.
116
+ * Can also be used as a codec
117
+ */
118
+ struct IndexLocalSearchQuantizerFastScan : IndexAdditiveQuantizerFastScan {
119
+ LocalSearchQuantizer lsq;
120
+
121
+ /** Constructor.
122
+ *
123
+ * @param d dimensionality of the input vectors
124
+ * @param M number of subquantizers
125
+ * @param nbits number of bit per subvector index
126
+ * @param metric metric type
127
+ * @param search_type AQ search type
128
+ */
129
+ IndexLocalSearchQuantizerFastScan(
130
+ int d, ///< dimensionality of the input vectors
131
+ size_t M, ///< number of subquantizers
132
+ size_t nbits, ///< number of bit per subvector index
133
+ MetricType metric = METRIC_L2,
134
+ Search_type_t search_type = AdditiveQuantizer::ST_norm_lsq2x4,
135
+ int bbs = 32);
136
+
137
+ IndexLocalSearchQuantizerFastScan();
138
+ };
139
+
140
+ /** Index based on a product residual quantizer. Stored vectors are
141
+ * approximated by product residual quantization codes.
142
+ * Can also be used as a codec
143
+ */
144
+ struct IndexProductResidualQuantizerFastScan : IndexAdditiveQuantizerFastScan {
145
+ /// The product residual quantizer used to encode the vectors
146
+ ProductResidualQuantizer prq;
147
+
148
+ /** Constructor.
149
+ *
150
+ * @param d dimensionality of the input vectors
151
+ * @param nsplits number of residual quantizers
152
+ * @param Msub number of subquantizers per RQ
153
+ * @param nbits number of bit per subvector index
154
+ * @param metric metric type
155
+ * @param search_type AQ search type
156
+ */
157
+ IndexProductResidualQuantizerFastScan(
158
+ int d, ///< dimensionality of the input vectors
159
+ size_t nsplits, ///< number of residual quantizers
160
+ size_t Msub, ///< number of subquantizers per RQ
161
+ size_t nbits, ///< number of bit per subvector index
162
+ MetricType metric = METRIC_L2,
163
+ Search_type_t search_type = AdditiveQuantizer::ST_norm_rq2x4,
164
+ int bbs = 32);
165
+
166
+ IndexProductResidualQuantizerFastScan();
167
+ };
168
+
169
+ /** Index based on a product local search quantizer. Stored vectors are
170
+ * approximated by product local search quantization codes.
171
+ * Can also be used as a codec
172
+ */
173
+ struct IndexProductLocalSearchQuantizerFastScan
174
+ : IndexAdditiveQuantizerFastScan {
175
+ /// The product local search quantizer used to encode the vectors
176
+ ProductLocalSearchQuantizer plsq;
177
+
178
+ /** Constructor.
179
+ *
180
+ * @param d dimensionality of the input vectors
181
+ * @param nsplits number of local search quantizers
182
+ * @param Msub number of subquantizers per LSQ
183
+ * @param nbits number of bit per subvector index
184
+ * @param metric metric type
185
+ * @param search_type AQ search type
186
+ */
187
+ IndexProductLocalSearchQuantizerFastScan(
188
+ int d, ///< dimensionality of the input vectors
189
+ size_t nsplits, ///< number of local search quantizers
190
+ size_t Msub, ///< number of subquantizers per LSQ
191
+ size_t nbits, ///< number of bit per subvector index
192
+ MetricType metric = METRIC_L2,
193
+ Search_type_t search_type = AdditiveQuantizer::ST_norm_rq2x4,
194
+ int bbs = 32);
195
+
196
+ IndexProductLocalSearchQuantizerFastScan();
197
+ };
198
+
199
+ } // namespace faiss
@@ -21,8 +21,12 @@ void IndexBinary::train(idx_t, const uint8_t*) {
21
21
  // Does nothing by default.
22
22
  }
23
23
 
24
- void IndexBinary::range_search(idx_t, const uint8_t*, int, RangeSearchResult*)
25
- const {
24
+ void IndexBinary::range_search(
25
+ idx_t,
26
+ const uint8_t*,
27
+ int,
28
+ RangeSearchResult*,
29
+ const SearchParameters*) const {
26
30
  FAISS_THROW_MSG("range search not implemented");
27
31
  }
28
32
 
@@ -57,10 +61,11 @@ void IndexBinary::search_and_reconstruct(
57
61
  idx_t k,
58
62
  int32_t* distances,
59
63
  idx_t* labels,
60
- uint8_t* recons) const {
64
+ uint8_t* recons,
65
+ const SearchParameters* params) const {
61
66
  FAISS_THROW_IF_NOT(k > 0);
62
67
 
63
- search(n, x, k, distances, labels);
68
+ search(n, x, k, distances, labels, params);
64
69
  for (idx_t i = 0; i < n; ++i) {
65
70
  for (idx_t j = 0; j < k; ++j) {
66
71
  idx_t ij = i * k + j;
@@ -82,4 +87,15 @@ void IndexBinary::display() const {
82
87
  ntotal);
83
88
  }
84
89
 
90
+ void IndexBinary::merge_from(
91
+ IndexBinary& /* otherIndex */,
92
+ idx_t /* add_id */) {
93
+ FAISS_THROW_MSG("merge_from() not implemented");
94
+ }
95
+
96
+ void IndexBinary::check_compatible_for_merge(
97
+ const IndexBinary& /* otherIndex */) const {
98
+ FAISS_THROW_MSG("check_compatible_for_merge() not implemented");
99
+ }
100
+
85
101
  } // namespace faiss
@@ -97,7 +97,8 @@ struct IndexBinary {
97
97
  const uint8_t* x,
98
98
  idx_t k,
99
99
  int32_t* distances,
100
- idx_t* labels) const = 0;
100
+ idx_t* labels,
101
+ const SearchParameters* params = nullptr) const = 0;
101
102
 
102
103
  /** Query n vectors of dimension d to the index.
103
104
  *
@@ -117,7 +118,8 @@ struct IndexBinary {
117
118
  idx_t n,
118
119
  const uint8_t* x,
119
120
  int radius,
120
- RangeSearchResult* result) const;
121
+ RangeSearchResult* result,
122
+ const SearchParameters* params = nullptr) const;
121
123
 
122
124
  /** Return the indexes of the k vectors closest to the query x.
123
125
  *
@@ -164,10 +166,23 @@ struct IndexBinary {
164
166
  idx_t k,
165
167
  int32_t* distances,
166
168
  idx_t* labels,
167
- uint8_t* recons) const;
169
+ uint8_t* recons,
170
+ const SearchParameters* params = nullptr) const;
168
171
 
169
172
  /** Display the actual class name and some more info. */
170
173
  void display() const;
174
+
175
+ /** moves the entries from another dataset to self.
176
+ * On output, other is empty.
177
+ * add_id is added to all moved ids
178
+ * (for sequential ids, this would be this->ntotal) */
179
+ virtual void merge_from(IndexBinary& otherIndex, idx_t add_id = 0);
180
+
181
+ /** check that the two indexes are compatible (ie, they are
182
+ * trained in the same way and have the same
183
+ * parameters). Otherwise throw. */
184
+ virtual void check_compatible_for_merge(
185
+ const IndexBinary& otherIndex) const;
171
186
  };
172
187
 
173
188
  } // namespace faiss
@@ -11,6 +11,7 @@
11
11
 
12
12
  #include <faiss/impl/AuxIndexStructures.h>
13
13
  #include <faiss/impl/FaissAssert.h>
14
+ #include <faiss/impl/IDSelector.h>
14
15
  #include <faiss/utils/Heap.h>
15
16
  #include <faiss/utils/hamming.h>
16
17
  #include <faiss/utils/utils.h>
@@ -35,7 +36,10 @@ void IndexBinaryFlat::search(
35
36
  const uint8_t* x,
36
37
  idx_t k,
37
38
  int32_t* distances,
38
- idx_t* labels) const {
39
+ idx_t* labels,
40
+ const SearchParameters* params) const {
41
+ FAISS_THROW_IF_NOT_MSG(
42
+ !params, "search params not supported for this index");
39
43
  FAISS_THROW_IF_NOT(k > 0);
40
44
 
41
45
  const idx_t block_size = query_batch_size;
@@ -101,7 +105,10 @@ void IndexBinaryFlat::range_search(
101
105
  idx_t n,
102
106
  const uint8_t* x,
103
107
  int radius,
104
- RangeSearchResult* result) const {
108
+ RangeSearchResult* result,
109
+ const SearchParameters* params) const {
110
+ FAISS_THROW_IF_NOT_MSG(
111
+ !params, "search params not supported for this index");
105
112
  hamming_range_search(x, xb.data(), n, ntotal, radius, code_size, result);
106
113
  }
107
114
 
@@ -39,13 +39,15 @@ struct IndexBinaryFlat : IndexBinary {
39
39
  const uint8_t* x,
40
40
  idx_t k,
41
41
  int32_t* distances,
42
- idx_t* labels) const override;
42
+ idx_t* labels,
43
+ const SearchParameters* params = nullptr) const override;
43
44
 
44
45
  void range_search(
45
46
  idx_t n,
46
47
  const uint8_t* x,
47
48
  int radius,
48
- RangeSearchResult* result) const override;
49
+ RangeSearchResult* result,
50
+ const SearchParameters* params = nullptr) const override;
49
51
 
50
52
  void reconstruct(idx_t key, uint8_t* recons) const override;
51
53
 
@@ -52,7 +52,10 @@ void IndexBinaryFromFloat::search(
52
52
  const uint8_t* x,
53
53
  idx_t k,
54
54
  int32_t* distances,
55
- idx_t* labels) const {
55
+ idx_t* labels,
56
+ const SearchParameters* params) const {
57
+ FAISS_THROW_IF_NOT_MSG(
58
+ !params, "search params not supported for this index");
56
59
  FAISS_THROW_IF_NOT(k > 0);
57
60
 
58
61
  constexpr idx_t bs = 32768;
@@ -43,7 +43,8 @@ struct IndexBinaryFromFloat : IndexBinary {
43
43
  const uint8_t* x,
44
44
  idx_t k,
45
45
  int32_t* distances,
46
- idx_t* labels) const override;
46
+ idx_t* labels,
47
+ const SearchParameters* params = nullptr) const override;
47
48
 
48
49
  void train(idx_t n, const uint8_t* x) override;
49
50
  };
@@ -26,6 +26,7 @@
26
26
 
27
27
  #include <faiss/IndexBinaryFlat.h>
28
28
  #include <faiss/impl/AuxIndexStructures.h>
29
+ #include <faiss/impl/DistanceComputer.h>
29
30
  #include <faiss/impl/FaissAssert.h>
30
31
  #include <faiss/utils/Heap.h>
31
32
  #include <faiss/utils/hamming.h>
@@ -194,7 +195,10 @@ void IndexBinaryHNSW::search(
194
195
  const uint8_t* x,
195
196
  idx_t k,
196
197
  int32_t* distances,
197
- idx_t* labels) const {
198
+ idx_t* labels,
199
+ const SearchParameters* params) const {
200
+ FAISS_THROW_IF_NOT_MSG(
201
+ !params, "search params not supported for this index");
198
202
  FAISS_THROW_IF_NOT(k > 0);
199
203
 
200
204
  #pragma omp parallel
@@ -47,7 +47,8 @@ struct IndexBinaryHNSW : IndexBinary {
47
47
  const uint8_t* x,
48
48
  idx_t k,
49
49
  int32_t* distances,
50
- idx_t* labels) const override;
50
+ idx_t* labels,
51
+ const SearchParameters* params = nullptr) const override;
51
52
 
52
53
  void reconstruct(idx_t key, uint8_t* recons) const override;
53
54