faiss 0.2.4 → 0.2.6
Sign up to get free protection for your applications and to get access to all the features.
- checksums.yaml +4 -4
- data/CHANGELOG.md +10 -0
- data/LICENSE.txt +1 -1
- data/README.md +23 -21
- data/ext/faiss/extconf.rb +11 -0
- data/ext/faiss/index.cpp +17 -4
- data/ext/faiss/index_binary.cpp +6 -6
- data/ext/faiss/product_quantizer.cpp +4 -4
- data/lib/faiss/version.rb +1 -1
- data/vendor/faiss/faiss/AutoTune.cpp +13 -0
- data/vendor/faiss/faiss/IVFlib.cpp +101 -2
- data/vendor/faiss/faiss/IVFlib.h +26 -2
- data/vendor/faiss/faiss/Index.cpp +36 -3
- data/vendor/faiss/faiss/Index.h +43 -6
- data/vendor/faiss/faiss/Index2Layer.cpp +6 -2
- data/vendor/faiss/faiss/Index2Layer.h +6 -1
- data/vendor/faiss/faiss/IndexAdditiveQuantizer.cpp +219 -16
- data/vendor/faiss/faiss/IndexAdditiveQuantizer.h +63 -5
- data/vendor/faiss/faiss/IndexAdditiveQuantizerFastScan.cpp +299 -0
- data/vendor/faiss/faiss/IndexAdditiveQuantizerFastScan.h +199 -0
- data/vendor/faiss/faiss/IndexBinary.cpp +20 -4
- data/vendor/faiss/faiss/IndexBinary.h +18 -3
- data/vendor/faiss/faiss/IndexBinaryFlat.cpp +9 -2
- data/vendor/faiss/faiss/IndexBinaryFlat.h +4 -2
- data/vendor/faiss/faiss/IndexBinaryFromFloat.cpp +4 -1
- data/vendor/faiss/faiss/IndexBinaryFromFloat.h +2 -1
- data/vendor/faiss/faiss/IndexBinaryHNSW.cpp +5 -1
- data/vendor/faiss/faiss/IndexBinaryHNSW.h +2 -1
- data/vendor/faiss/faiss/IndexBinaryHash.cpp +17 -4
- data/vendor/faiss/faiss/IndexBinaryHash.h +8 -4
- data/vendor/faiss/faiss/IndexBinaryIVF.cpp +28 -13
- data/vendor/faiss/faiss/IndexBinaryIVF.h +10 -7
- data/vendor/faiss/faiss/IndexFastScan.cpp +626 -0
- data/vendor/faiss/faiss/IndexFastScan.h +145 -0
- data/vendor/faiss/faiss/IndexFlat.cpp +34 -21
- data/vendor/faiss/faiss/IndexFlat.h +7 -4
- data/vendor/faiss/faiss/IndexFlatCodes.cpp +35 -1
- data/vendor/faiss/faiss/IndexFlatCodes.h +12 -0
- data/vendor/faiss/faiss/IndexHNSW.cpp +66 -138
- data/vendor/faiss/faiss/IndexHNSW.h +4 -2
- data/vendor/faiss/faiss/IndexIDMap.cpp +247 -0
- data/vendor/faiss/faiss/IndexIDMap.h +107 -0
- data/vendor/faiss/faiss/IndexIVF.cpp +121 -33
- data/vendor/faiss/faiss/IndexIVF.h +35 -16
- data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.cpp +84 -7
- data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.h +63 -1
- data/vendor/faiss/faiss/IndexIVFAdditiveQuantizerFastScan.cpp +590 -0
- data/vendor/faiss/faiss/IndexIVFAdditiveQuantizerFastScan.h +171 -0
- data/vendor/faiss/faiss/IndexIVFFastScan.cpp +1290 -0
- data/vendor/faiss/faiss/IndexIVFFastScan.h +213 -0
- data/vendor/faiss/faiss/IndexIVFFlat.cpp +37 -17
- data/vendor/faiss/faiss/IndexIVFFlat.h +4 -2
- data/vendor/faiss/faiss/IndexIVFPQ.cpp +234 -50
- data/vendor/faiss/faiss/IndexIVFPQ.h +5 -1
- data/vendor/faiss/faiss/IndexIVFPQFastScan.cpp +23 -852
- data/vendor/faiss/faiss/IndexIVFPQFastScan.h +7 -112
- data/vendor/faiss/faiss/IndexIVFPQR.cpp +3 -3
- data/vendor/faiss/faiss/IndexIVFPQR.h +1 -1
- data/vendor/faiss/faiss/IndexIVFSpectralHash.cpp +3 -1
- data/vendor/faiss/faiss/IndexIVFSpectralHash.h +2 -1
- data/vendor/faiss/faiss/IndexLSH.cpp +4 -2
- data/vendor/faiss/faiss/IndexLSH.h +2 -1
- data/vendor/faiss/faiss/IndexLattice.cpp +7 -1
- data/vendor/faiss/faiss/IndexLattice.h +3 -1
- data/vendor/faiss/faiss/IndexNNDescent.cpp +4 -3
- data/vendor/faiss/faiss/IndexNNDescent.h +2 -1
- data/vendor/faiss/faiss/IndexNSG.cpp +37 -3
- data/vendor/faiss/faiss/IndexNSG.h +25 -1
- data/vendor/faiss/faiss/IndexPQ.cpp +106 -69
- data/vendor/faiss/faiss/IndexPQ.h +19 -5
- data/vendor/faiss/faiss/IndexPQFastScan.cpp +15 -450
- data/vendor/faiss/faiss/IndexPQFastScan.h +15 -78
- data/vendor/faiss/faiss/IndexPreTransform.cpp +47 -8
- data/vendor/faiss/faiss/IndexPreTransform.h +15 -3
- data/vendor/faiss/faiss/IndexRefine.cpp +8 -4
- data/vendor/faiss/faiss/IndexRefine.h +4 -2
- data/vendor/faiss/faiss/IndexReplicas.cpp +4 -2
- data/vendor/faiss/faiss/IndexReplicas.h +2 -1
- data/vendor/faiss/faiss/IndexRowwiseMinMax.cpp +438 -0
- data/vendor/faiss/faiss/IndexRowwiseMinMax.h +92 -0
- data/vendor/faiss/faiss/IndexScalarQuantizer.cpp +26 -15
- data/vendor/faiss/faiss/IndexScalarQuantizer.h +6 -7
- data/vendor/faiss/faiss/IndexShards.cpp +4 -1
- data/vendor/faiss/faiss/IndexShards.h +2 -1
- data/vendor/faiss/faiss/MetaIndexes.cpp +5 -178
- data/vendor/faiss/faiss/MetaIndexes.h +3 -81
- data/vendor/faiss/faiss/VectorTransform.cpp +43 -0
- data/vendor/faiss/faiss/VectorTransform.h +22 -4
- data/vendor/faiss/faiss/clone_index.cpp +23 -1
- data/vendor/faiss/faiss/clone_index.h +3 -0
- data/vendor/faiss/faiss/cppcontrib/SaDecodeKernels.h +300 -0
- data/vendor/faiss/faiss/cppcontrib/detail/CoarseBitType.h +24 -0
- data/vendor/faiss/faiss/cppcontrib/detail/UintReader.h +195 -0
- data/vendor/faiss/faiss/cppcontrib/sa_decode/Level2-avx2-inl.h +2058 -0
- data/vendor/faiss/faiss/cppcontrib/sa_decode/Level2-inl.h +408 -0
- data/vendor/faiss/faiss/cppcontrib/sa_decode/Level2-neon-inl.h +2147 -0
- data/vendor/faiss/faiss/cppcontrib/sa_decode/MinMax-inl.h +460 -0
- data/vendor/faiss/faiss/cppcontrib/sa_decode/MinMaxFP16-inl.h +465 -0
- data/vendor/faiss/faiss/cppcontrib/sa_decode/PQ-avx2-inl.h +1618 -0
- data/vendor/faiss/faiss/cppcontrib/sa_decode/PQ-inl.h +251 -0
- data/vendor/faiss/faiss/cppcontrib/sa_decode/PQ-neon-inl.h +1452 -0
- data/vendor/faiss/faiss/gpu/GpuAutoTune.cpp +1 -0
- data/vendor/faiss/faiss/gpu/GpuCloner.cpp +0 -4
- data/vendor/faiss/faiss/gpu/GpuIndex.h +28 -4
- data/vendor/faiss/faiss/gpu/GpuIndexBinaryFlat.h +2 -1
- data/vendor/faiss/faiss/gpu/GpuIndexFlat.h +10 -8
- data/vendor/faiss/faiss/gpu/GpuIndexIVF.h +75 -14
- data/vendor/faiss/faiss/gpu/GpuIndexIVFFlat.h +19 -32
- data/vendor/faiss/faiss/gpu/GpuIndexIVFPQ.h +22 -31
- data/vendor/faiss/faiss/gpu/GpuIndexIVFScalarQuantizer.h +22 -28
- data/vendor/faiss/faiss/gpu/GpuResources.cpp +14 -0
- data/vendor/faiss/faiss/gpu/GpuResources.h +16 -3
- data/vendor/faiss/faiss/gpu/StandardGpuResources.cpp +3 -3
- data/vendor/faiss/faiss/gpu/impl/IndexUtils.h +32 -0
- data/vendor/faiss/faiss/gpu/test/TestGpuIndexBinaryFlat.cpp +1 -0
- data/vendor/faiss/faiss/gpu/test/TestGpuIndexFlat.cpp +311 -75
- data/vendor/faiss/faiss/gpu/test/TestUtils.cpp +10 -0
- data/vendor/faiss/faiss/gpu/test/TestUtils.h +3 -0
- data/vendor/faiss/faiss/gpu/test/demo_ivfpq_indexing_gpu.cpp +2 -2
- data/vendor/faiss/faiss/gpu/utils/DeviceUtils.h +5 -4
- data/vendor/faiss/faiss/impl/AdditiveQuantizer.cpp +116 -47
- data/vendor/faiss/faiss/impl/AdditiveQuantizer.h +44 -13
- data/vendor/faiss/faiss/impl/AuxIndexStructures.cpp +0 -54
- data/vendor/faiss/faiss/impl/AuxIndexStructures.h +0 -76
- data/vendor/faiss/faiss/impl/DistanceComputer.h +64 -0
- data/vendor/faiss/faiss/impl/HNSW.cpp +123 -27
- data/vendor/faiss/faiss/impl/HNSW.h +19 -16
- data/vendor/faiss/faiss/impl/IDSelector.cpp +125 -0
- data/vendor/faiss/faiss/impl/IDSelector.h +135 -0
- data/vendor/faiss/faiss/impl/LocalSearchQuantizer.cpp +6 -28
- data/vendor/faiss/faiss/impl/LocalSearchQuantizer.h +6 -1
- data/vendor/faiss/faiss/impl/LookupTableScaler.h +77 -0
- data/vendor/faiss/faiss/impl/NNDescent.cpp +1 -0
- data/vendor/faiss/faiss/impl/NSG.cpp +1 -1
- data/vendor/faiss/faiss/impl/ProductAdditiveQuantizer.cpp +383 -0
- data/vendor/faiss/faiss/impl/ProductAdditiveQuantizer.h +154 -0
- data/vendor/faiss/faiss/impl/ProductQuantizer.cpp +225 -145
- data/vendor/faiss/faiss/impl/ProductQuantizer.h +29 -10
- data/vendor/faiss/faiss/impl/Quantizer.h +43 -0
- data/vendor/faiss/faiss/impl/ResidualQuantizer.cpp +192 -36
- data/vendor/faiss/faiss/impl/ResidualQuantizer.h +40 -20
- data/vendor/faiss/faiss/impl/ResultHandler.h +96 -0
- data/vendor/faiss/faiss/impl/ScalarQuantizer.cpp +97 -173
- data/vendor/faiss/faiss/impl/ScalarQuantizer.h +18 -18
- data/vendor/faiss/faiss/impl/index_read.cpp +240 -9
- data/vendor/faiss/faiss/impl/index_write.cpp +237 -5
- data/vendor/faiss/faiss/impl/kmeans1d.cpp +6 -4
- data/vendor/faiss/faiss/impl/pq4_fast_scan.cpp +56 -16
- data/vendor/faiss/faiss/impl/pq4_fast_scan.h +25 -8
- data/vendor/faiss/faiss/impl/pq4_fast_scan_search_1.cpp +66 -25
- data/vendor/faiss/faiss/impl/pq4_fast_scan_search_qbs.cpp +75 -27
- data/vendor/faiss/faiss/index_factory.cpp +196 -7
- data/vendor/faiss/faiss/index_io.h +5 -0
- data/vendor/faiss/faiss/invlists/DirectMap.cpp +1 -0
- data/vendor/faiss/faiss/invlists/InvertedLists.cpp +4 -1
- data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.cpp +2 -1
- data/vendor/faiss/faiss/python/python_callbacks.cpp +27 -0
- data/vendor/faiss/faiss/python/python_callbacks.h +15 -0
- data/vendor/faiss/faiss/utils/Heap.h +31 -15
- data/vendor/faiss/faiss/utils/distances.cpp +380 -56
- data/vendor/faiss/faiss/utils/distances.h +113 -15
- data/vendor/faiss/faiss/utils/distances_simd.cpp +726 -6
- data/vendor/faiss/faiss/utils/extra_distances.cpp +12 -7
- data/vendor/faiss/faiss/utils/extra_distances.h +3 -1
- data/vendor/faiss/faiss/utils/fp16-fp16c.h +21 -0
- data/vendor/faiss/faiss/utils/fp16-inl.h +101 -0
- data/vendor/faiss/faiss/utils/fp16.h +11 -0
- data/vendor/faiss/faiss/utils/hamming-inl.h +54 -0
- data/vendor/faiss/faiss/utils/hamming.cpp +0 -48
- data/vendor/faiss/faiss/utils/ordered_key_value.h +10 -0
- data/vendor/faiss/faiss/utils/quantize_lut.cpp +62 -0
- data/vendor/faiss/faiss/utils/quantize_lut.h +20 -0
- data/vendor/faiss/faiss/utils/random.cpp +53 -0
- data/vendor/faiss/faiss/utils/random.h +5 -0
- data/vendor/faiss/faiss/utils/simdlib_avx2.h +4 -0
- data/vendor/faiss/faiss/utils/simdlib_emulated.h +6 -1
- data/vendor/faiss/faiss/utils/simdlib_neon.h +7 -2
- 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(
|
25
|
-
|
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
|
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
|
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
|
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
|
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
|
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
|
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
|
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
|
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
|
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
|
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
|
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
|
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
|
|