faiss 0.2.6 → 0.2.7
Sign up to get free protection for your applications and to get access to all the features.
- checksums.yaml +4 -4
- data/CHANGELOG.md +4 -0
- data/ext/faiss/extconf.rb +1 -1
- data/lib/faiss/version.rb +1 -1
- data/lib/faiss.rb +2 -2
- data/vendor/faiss/faiss/AutoTune.cpp +15 -4
- data/vendor/faiss/faiss/AutoTune.h +0 -1
- data/vendor/faiss/faiss/Clustering.cpp +1 -5
- data/vendor/faiss/faiss/Clustering.h +0 -2
- data/vendor/faiss/faiss/IVFlib.h +0 -2
- data/vendor/faiss/faiss/Index.h +1 -2
- data/vendor/faiss/faiss/IndexAdditiveQuantizer.cpp +17 -3
- data/vendor/faiss/faiss/IndexAdditiveQuantizer.h +10 -1
- data/vendor/faiss/faiss/IndexBinary.h +0 -1
- data/vendor/faiss/faiss/IndexBinaryFlat.cpp +2 -1
- data/vendor/faiss/faiss/IndexBinaryFlat.h +4 -0
- data/vendor/faiss/faiss/IndexBinaryHash.cpp +1 -3
- data/vendor/faiss/faiss/IndexBinaryIVF.cpp +273 -48
- data/vendor/faiss/faiss/IndexBinaryIVF.h +18 -11
- data/vendor/faiss/faiss/IndexFastScan.cpp +13 -10
- data/vendor/faiss/faiss/IndexFastScan.h +5 -1
- data/vendor/faiss/faiss/IndexFlat.cpp +16 -3
- data/vendor/faiss/faiss/IndexFlat.h +1 -1
- data/vendor/faiss/faiss/IndexFlatCodes.cpp +5 -0
- data/vendor/faiss/faiss/IndexFlatCodes.h +7 -2
- data/vendor/faiss/faiss/IndexHNSW.cpp +3 -6
- data/vendor/faiss/faiss/IndexHNSW.h +0 -1
- data/vendor/faiss/faiss/IndexIDMap.cpp +4 -4
- data/vendor/faiss/faiss/IndexIDMap.h +0 -2
- data/vendor/faiss/faiss/IndexIVF.cpp +155 -129
- data/vendor/faiss/faiss/IndexIVF.h +121 -61
- data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.cpp +2 -2
- data/vendor/faiss/faiss/IndexIVFFastScan.cpp +12 -11
- data/vendor/faiss/faiss/IndexIVFFastScan.h +6 -1
- data/vendor/faiss/faiss/IndexIVFPQ.cpp +221 -165
- data/vendor/faiss/faiss/IndexIVFPQ.h +1 -0
- data/vendor/faiss/faiss/IndexIVFPQFastScan.cpp +6 -1
- data/vendor/faiss/faiss/IndexIVFSpectralHash.cpp +0 -2
- data/vendor/faiss/faiss/IndexNNDescent.cpp +1 -2
- data/vendor/faiss/faiss/IndexNNDescent.h +0 -1
- data/vendor/faiss/faiss/IndexNSG.cpp +1 -2
- data/vendor/faiss/faiss/IndexPQ.cpp +7 -9
- data/vendor/faiss/faiss/IndexRefine.cpp +1 -1
- data/vendor/faiss/faiss/IndexReplicas.cpp +3 -4
- data/vendor/faiss/faiss/IndexReplicas.h +0 -1
- data/vendor/faiss/faiss/IndexRowwiseMinMax.cpp +8 -1
- data/vendor/faiss/faiss/IndexRowwiseMinMax.h +7 -0
- data/vendor/faiss/faiss/IndexShards.cpp +26 -109
- data/vendor/faiss/faiss/IndexShards.h +2 -3
- data/vendor/faiss/faiss/IndexShardsIVF.cpp +246 -0
- data/vendor/faiss/faiss/IndexShardsIVF.h +42 -0
- data/vendor/faiss/faiss/MetaIndexes.cpp +86 -0
- data/vendor/faiss/faiss/MetaIndexes.h +29 -0
- data/vendor/faiss/faiss/MetricType.h +14 -0
- data/vendor/faiss/faiss/VectorTransform.cpp +8 -10
- data/vendor/faiss/faiss/VectorTransform.h +1 -3
- data/vendor/faiss/faiss/clone_index.cpp +232 -18
- data/vendor/faiss/faiss/cppcontrib/SaDecodeKernels.h +25 -3
- data/vendor/faiss/faiss/cppcontrib/detail/CoarseBitType.h +7 -0
- data/vendor/faiss/faiss/cppcontrib/detail/UintReader.h +78 -0
- data/vendor/faiss/faiss/cppcontrib/sa_decode/Level2-avx2-inl.h +20 -6
- data/vendor/faiss/faiss/cppcontrib/sa_decode/Level2-inl.h +7 -1
- data/vendor/faiss/faiss/cppcontrib/sa_decode/Level2-neon-inl.h +21 -7
- data/vendor/faiss/faiss/cppcontrib/sa_decode/MinMax-inl.h +7 -0
- data/vendor/faiss/faiss/cppcontrib/sa_decode/MinMaxFP16-inl.h +7 -0
- data/vendor/faiss/faiss/cppcontrib/sa_decode/PQ-avx2-inl.h +10 -3
- data/vendor/faiss/faiss/cppcontrib/sa_decode/PQ-inl.h +7 -1
- data/vendor/faiss/faiss/cppcontrib/sa_decode/PQ-neon-inl.h +11 -3
- data/vendor/faiss/faiss/gpu/GpuAutoTune.cpp +25 -2
- data/vendor/faiss/faiss/gpu/GpuCloner.cpp +76 -29
- data/vendor/faiss/faiss/gpu/GpuCloner.h +2 -2
- data/vendor/faiss/faiss/gpu/GpuClonerOptions.h +14 -13
- data/vendor/faiss/faiss/gpu/GpuDistance.h +18 -6
- data/vendor/faiss/faiss/gpu/GpuIndex.h +23 -21
- data/vendor/faiss/faiss/gpu/GpuIndexBinaryFlat.h +10 -10
- data/vendor/faiss/faiss/gpu/GpuIndexFlat.h +11 -12
- data/vendor/faiss/faiss/gpu/GpuIndexIVF.h +29 -50
- data/vendor/faiss/faiss/gpu/GpuIndexIVFFlat.h +3 -3
- data/vendor/faiss/faiss/gpu/GpuIndexIVFPQ.h +8 -8
- data/vendor/faiss/faiss/gpu/GpuIndexIVFScalarQuantizer.h +4 -4
- data/vendor/faiss/faiss/gpu/impl/IndexUtils.h +2 -5
- data/vendor/faiss/faiss/gpu/impl/RemapIndices.cpp +9 -7
- data/vendor/faiss/faiss/gpu/impl/RemapIndices.h +4 -4
- data/vendor/faiss/faiss/gpu/perf/IndexWrapper-inl.h +2 -2
- data/vendor/faiss/faiss/gpu/perf/IndexWrapper.h +1 -1
- data/vendor/faiss/faiss/gpu/test/TestGpuIndexBinaryFlat.cpp +55 -6
- data/vendor/faiss/faiss/gpu/test/TestGpuIndexFlat.cpp +20 -6
- data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFFlat.cpp +95 -25
- data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFPQ.cpp +67 -16
- data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFScalarQuantizer.cpp +4 -4
- data/vendor/faiss/faiss/gpu/test/TestUtils.cpp +7 -7
- data/vendor/faiss/faiss/gpu/test/TestUtils.h +4 -4
- data/vendor/faiss/faiss/gpu/test/demo_ivfpq_indexing_gpu.cpp +1 -1
- data/vendor/faiss/faiss/gpu/utils/DeviceUtils.h +6 -0
- data/vendor/faiss/faiss/impl/AdditiveQuantizer.cpp +0 -7
- data/vendor/faiss/faiss/impl/AdditiveQuantizer.h +9 -9
- data/vendor/faiss/faiss/impl/AuxIndexStructures.cpp +1 -1
- data/vendor/faiss/faiss/impl/AuxIndexStructures.h +2 -7
- data/vendor/faiss/faiss/impl/CodePacker.cpp +67 -0
- data/vendor/faiss/faiss/impl/CodePacker.h +71 -0
- data/vendor/faiss/faiss/impl/DistanceComputer.h +0 -2
- data/vendor/faiss/faiss/impl/HNSW.cpp +3 -7
- data/vendor/faiss/faiss/impl/HNSW.h +6 -9
- data/vendor/faiss/faiss/impl/IDSelector.cpp +1 -1
- data/vendor/faiss/faiss/impl/IDSelector.h +39 -1
- data/vendor/faiss/faiss/impl/LocalSearchQuantizer.cpp +62 -51
- data/vendor/faiss/faiss/impl/LocalSearchQuantizer.h +11 -12
- data/vendor/faiss/faiss/impl/NNDescent.cpp +3 -9
- data/vendor/faiss/faiss/impl/NNDescent.h +10 -10
- data/vendor/faiss/faiss/impl/NSG.cpp +1 -6
- data/vendor/faiss/faiss/impl/NSG.h +4 -7
- data/vendor/faiss/faiss/impl/PolysemousTraining.cpp +1 -15
- data/vendor/faiss/faiss/impl/PolysemousTraining.h +11 -10
- data/vendor/faiss/faiss/impl/ProductAdditiveQuantizer.cpp +0 -7
- data/vendor/faiss/faiss/impl/ProductQuantizer.cpp +25 -12
- data/vendor/faiss/faiss/impl/ProductQuantizer.h +2 -4
- data/vendor/faiss/faiss/impl/Quantizer.h +6 -3
- data/vendor/faiss/faiss/impl/ResidualQuantizer.cpp +796 -174
- data/vendor/faiss/faiss/impl/ResidualQuantizer.h +16 -8
- data/vendor/faiss/faiss/impl/ScalarQuantizer.cpp +3 -5
- data/vendor/faiss/faiss/impl/ScalarQuantizer.h +4 -4
- data/vendor/faiss/faiss/impl/ThreadedIndex-inl.h +3 -3
- data/vendor/faiss/faiss/impl/ThreadedIndex.h +4 -4
- data/vendor/faiss/faiss/impl/code_distance/code_distance-avx2.h +291 -0
- data/vendor/faiss/faiss/impl/code_distance/code_distance-generic.h +74 -0
- data/vendor/faiss/faiss/impl/code_distance/code_distance.h +123 -0
- data/vendor/faiss/faiss/impl/code_distance/code_distance_avx512.h +102 -0
- data/vendor/faiss/faiss/impl/index_read.cpp +13 -10
- data/vendor/faiss/faiss/impl/index_write.cpp +3 -4
- data/vendor/faiss/faiss/impl/kmeans1d.cpp +0 -1
- data/vendor/faiss/faiss/impl/kmeans1d.h +3 -3
- data/vendor/faiss/faiss/impl/lattice_Zn.cpp +1 -1
- data/vendor/faiss/faiss/impl/platform_macros.h +61 -0
- data/vendor/faiss/faiss/impl/pq4_fast_scan.cpp +48 -4
- data/vendor/faiss/faiss/impl/pq4_fast_scan.h +18 -4
- data/vendor/faiss/faiss/impl/pq4_fast_scan_search_qbs.cpp +2 -2
- data/vendor/faiss/faiss/index_factory.cpp +8 -10
- data/vendor/faiss/faiss/invlists/BlockInvertedLists.cpp +29 -12
- data/vendor/faiss/faiss/invlists/BlockInvertedLists.h +8 -2
- data/vendor/faiss/faiss/invlists/DirectMap.cpp +1 -1
- data/vendor/faiss/faiss/invlists/DirectMap.h +2 -4
- data/vendor/faiss/faiss/invlists/InvertedLists.cpp +118 -18
- data/vendor/faiss/faiss/invlists/InvertedLists.h +44 -4
- data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.cpp +3 -3
- data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.h +1 -1
- data/vendor/faiss/faiss/python/python_callbacks.cpp +1 -1
- data/vendor/faiss/faiss/python/python_callbacks.h +1 -1
- data/vendor/faiss/faiss/utils/AlignedTable.h +3 -1
- data/vendor/faiss/faiss/utils/Heap.cpp +139 -3
- data/vendor/faiss/faiss/utils/Heap.h +35 -1
- data/vendor/faiss/faiss/utils/approx_topk/approx_topk.h +84 -0
- data/vendor/faiss/faiss/utils/approx_topk/avx2-inl.h +196 -0
- data/vendor/faiss/faiss/utils/approx_topk/generic.h +138 -0
- data/vendor/faiss/faiss/utils/approx_topk/mode.h +34 -0
- data/vendor/faiss/faiss/utils/approx_topk_hamming/approx_topk_hamming.h +367 -0
- data/vendor/faiss/faiss/utils/distances.cpp +61 -7
- data/vendor/faiss/faiss/utils/distances.h +11 -0
- data/vendor/faiss/faiss/utils/distances_fused/avx512.cpp +346 -0
- data/vendor/faiss/faiss/utils/distances_fused/avx512.h +36 -0
- data/vendor/faiss/faiss/utils/distances_fused/distances_fused.cpp +42 -0
- data/vendor/faiss/faiss/utils/distances_fused/distances_fused.h +40 -0
- data/vendor/faiss/faiss/utils/distances_fused/simdlib_based.cpp +352 -0
- data/vendor/faiss/faiss/utils/distances_fused/simdlib_based.h +32 -0
- data/vendor/faiss/faiss/utils/distances_simd.cpp +515 -327
- data/vendor/faiss/faiss/utils/extra_distances-inl.h +17 -1
- data/vendor/faiss/faiss/utils/extra_distances.cpp +37 -8
- data/vendor/faiss/faiss/utils/extra_distances.h +2 -1
- data/vendor/faiss/faiss/utils/fp16-fp16c.h +7 -0
- data/vendor/faiss/faiss/utils/fp16-inl.h +7 -0
- data/vendor/faiss/faiss/utils/fp16.h +7 -0
- data/vendor/faiss/faiss/utils/hamming-inl.h +0 -456
- data/vendor/faiss/faiss/utils/hamming.cpp +104 -120
- data/vendor/faiss/faiss/utils/hamming.h +21 -10
- data/vendor/faiss/faiss/utils/hamming_distance/avx2-inl.h +535 -0
- data/vendor/faiss/faiss/utils/hamming_distance/common.h +48 -0
- data/vendor/faiss/faiss/utils/hamming_distance/generic-inl.h +519 -0
- data/vendor/faiss/faiss/utils/hamming_distance/hamdis-inl.h +26 -0
- data/vendor/faiss/faiss/utils/hamming_distance/neon-inl.h +614 -0
- data/vendor/faiss/faiss/utils/partitioning.cpp +21 -25
- data/vendor/faiss/faiss/utils/simdlib_avx2.h +344 -3
- data/vendor/faiss/faiss/utils/simdlib_emulated.h +390 -0
- data/vendor/faiss/faiss/utils/simdlib_neon.h +655 -130
- data/vendor/faiss/faiss/utils/sorting.cpp +692 -0
- data/vendor/faiss/faiss/utils/sorting.h +71 -0
- data/vendor/faiss/faiss/utils/transpose/transpose-avx2-inl.h +165 -0
- data/vendor/faiss/faiss/utils/utils.cpp +4 -176
- data/vendor/faiss/faiss/utils/utils.h +2 -9
- metadata +29 -3
- data/vendor/faiss/faiss/gpu/GpuClonerOptions.cpp +0 -26
@@ -64,13 +64,10 @@ struct Options {
|
|
64
64
|
};
|
65
65
|
|
66
66
|
void queryTest(
|
67
|
+
Options opt,
|
67
68
|
faiss::MetricType metricType,
|
68
|
-
bool useFloat16CoarseQuantizer
|
69
|
-
int dimOverride = -1) {
|
69
|
+
bool useFloat16CoarseQuantizer) {
|
70
70
|
for (int tries = 0; tries < 2; ++tries) {
|
71
|
-
Options opt;
|
72
|
-
opt.dim = dimOverride != -1 ? dimOverride : opt.dim;
|
73
|
-
|
74
71
|
std::vector<float> trainVecs =
|
75
72
|
faiss::gpu::randVecs(opt.numTrain, opt.dim);
|
76
73
|
std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
|
@@ -98,7 +95,7 @@ void queryTest(
|
|
98
95
|
faiss::gpu::GpuIndexIVFFlat gpuIndex(
|
99
96
|
&res, cpuIndex.d, cpuIndex.nlist, cpuIndex.metric_type, config);
|
100
97
|
gpuIndex.copyFrom(&cpuIndex);
|
101
|
-
gpuIndex.
|
98
|
+
gpuIndex.nprobe = opt.nprobe;
|
102
99
|
|
103
100
|
bool compFloat16 = useFloat16CoarseQuantizer;
|
104
101
|
faiss::gpu::compareIndices(
|
@@ -147,7 +144,7 @@ void addTest(faiss::MetricType metricType, bool useFloat16CoarseQuantizer) {
|
|
147
144
|
faiss::gpu::GpuIndexIVFFlat gpuIndex(
|
148
145
|
&res, cpuIndex.d, cpuIndex.nlist, cpuIndex.metric_type, config);
|
149
146
|
gpuIndex.copyFrom(&cpuIndex);
|
150
|
-
gpuIndex.
|
147
|
+
gpuIndex.nprobe = opt.nprobe;
|
151
148
|
|
152
149
|
cpuIndex.add(opt.numAdd, addVecs.data());
|
153
150
|
gpuIndex.add(opt.numAdd, addVecs.data());
|
@@ -183,7 +180,7 @@ void copyToTest(bool useFloat16CoarseQuantizer) {
|
|
183
180
|
&res, opt.dim, opt.numCentroids, faiss::METRIC_L2, config);
|
184
181
|
gpuIndex.train(opt.numTrain, trainVecs.data());
|
185
182
|
gpuIndex.add(opt.numAdd, addVecs.data());
|
186
|
-
gpuIndex.
|
183
|
+
gpuIndex.nprobe = opt.nprobe;
|
187
184
|
|
188
185
|
// use garbage values to see if we overwrite then
|
189
186
|
faiss::IndexFlatL2 cpuQuantizer(1);
|
@@ -199,7 +196,7 @@ void copyToTest(bool useFloat16CoarseQuantizer) {
|
|
199
196
|
EXPECT_EQ(cpuIndex.quantizer->d, gpuIndex.quantizer->d);
|
200
197
|
EXPECT_EQ(cpuIndex.d, opt.dim);
|
201
198
|
EXPECT_EQ(cpuIndex.nlist, gpuIndex.getNumLists());
|
202
|
-
EXPECT_EQ(cpuIndex.nprobe, gpuIndex.
|
199
|
+
EXPECT_EQ(cpuIndex.nprobe, gpuIndex.nprobe);
|
203
200
|
|
204
201
|
testIVFEquality(cpuIndex, gpuIndex);
|
205
202
|
|
@@ -239,7 +236,7 @@ void copyFromTest(bool useFloat16CoarseQuantizer) {
|
|
239
236
|
config.flatConfig.useFloat16 = useFloat16CoarseQuantizer;
|
240
237
|
|
241
238
|
faiss::gpu::GpuIndexIVFFlat gpuIndex(&res, 1, 1, faiss::METRIC_L2, config);
|
242
|
-
gpuIndex.
|
239
|
+
gpuIndex.nprobe = 1;
|
243
240
|
|
244
241
|
gpuIndex.copyFrom(&cpuIndex);
|
245
242
|
|
@@ -249,7 +246,7 @@ void copyFromTest(bool useFloat16CoarseQuantizer) {
|
|
249
246
|
EXPECT_EQ(cpuIndex.d, gpuIndex.d);
|
250
247
|
EXPECT_EQ(cpuIndex.d, opt.dim);
|
251
248
|
EXPECT_EQ(cpuIndex.nlist, gpuIndex.getNumLists());
|
252
|
-
EXPECT_EQ(cpuIndex.nprobe, gpuIndex.
|
249
|
+
EXPECT_EQ(cpuIndex.nprobe, gpuIndex.nprobe);
|
253
250
|
|
254
251
|
testIVFEquality(cpuIndex, gpuIndex);
|
255
252
|
|
@@ -288,21 +285,28 @@ TEST(TestGpuIndexIVFFlat, Float16_32_Add_IP) {
|
|
288
285
|
//
|
289
286
|
|
290
287
|
TEST(TestGpuIndexIVFFlat, Float32_Query_L2) {
|
291
|
-
queryTest(faiss::METRIC_L2, false);
|
288
|
+
queryTest(Options(), faiss::METRIC_L2, false);
|
292
289
|
}
|
293
290
|
|
294
291
|
TEST(TestGpuIndexIVFFlat, Float32_Query_IP) {
|
295
|
-
queryTest(faiss::METRIC_INNER_PRODUCT, false);
|
292
|
+
queryTest(Options(), faiss::METRIC_INNER_PRODUCT, false);
|
293
|
+
}
|
294
|
+
|
295
|
+
TEST(TestGpuIndexIVFFlat, LargeBatch) {
|
296
|
+
Options opt;
|
297
|
+
opt.dim = 3;
|
298
|
+
opt.numQuery = 100000;
|
299
|
+
queryTest(opt, faiss::METRIC_L2, false);
|
296
300
|
}
|
297
301
|
|
298
302
|
// float16 coarse quantizer
|
299
303
|
|
300
304
|
TEST(TestGpuIndexIVFFlat, Float16_32_Query_L2) {
|
301
|
-
queryTest(faiss::METRIC_L2, true);
|
305
|
+
queryTest(Options(), faiss::METRIC_L2, true);
|
302
306
|
}
|
303
307
|
|
304
308
|
TEST(TestGpuIndexIVFFlat, Float16_32_Query_IP) {
|
305
|
-
queryTest(faiss::METRIC_INNER_PRODUCT, true);
|
309
|
+
queryTest(Options(), faiss::METRIC_INNER_PRODUCT, true);
|
306
310
|
}
|
307
311
|
|
308
312
|
//
|
@@ -311,19 +315,27 @@ TEST(TestGpuIndexIVFFlat, Float16_32_Query_IP) {
|
|
311
315
|
//
|
312
316
|
|
313
317
|
TEST(TestGpuIndexIVFFlat, Float32_Query_L2_64) {
|
314
|
-
|
318
|
+
Options opt;
|
319
|
+
opt.dim = 64;
|
320
|
+
queryTest(opt, faiss::METRIC_L2, false);
|
315
321
|
}
|
316
322
|
|
317
323
|
TEST(TestGpuIndexIVFFlat, Float32_Query_IP_64) {
|
318
|
-
|
324
|
+
Options opt;
|
325
|
+
opt.dim = 64;
|
326
|
+
queryTest(opt, faiss::METRIC_INNER_PRODUCT, false);
|
319
327
|
}
|
320
328
|
|
321
329
|
TEST(TestGpuIndexIVFFlat, Float32_Query_L2_128) {
|
322
|
-
|
330
|
+
Options opt;
|
331
|
+
opt.dim = 128;
|
332
|
+
queryTest(opt, faiss::METRIC_L2, false);
|
323
333
|
}
|
324
334
|
|
325
335
|
TEST(TestGpuIndexIVFFlat, Float32_Query_IP_128) {
|
326
|
-
|
336
|
+
Options opt;
|
337
|
+
opt.dim = 128;
|
338
|
+
queryTest(opt, faiss::METRIC_INNER_PRODUCT, false);
|
327
339
|
}
|
328
340
|
|
329
341
|
//
|
@@ -372,7 +384,7 @@ TEST(TestGpuIndexIVFFlat, Float32_negative) {
|
|
372
384
|
faiss::gpu::GpuIndexIVFFlat gpuIndex(
|
373
385
|
&res, cpuIndex.d, cpuIndex.nlist, cpuIndex.metric_type, config);
|
374
386
|
gpuIndex.copyFrom(&cpuIndex);
|
375
|
-
gpuIndex.
|
387
|
+
gpuIndex.nprobe = opt.nprobe;
|
376
388
|
|
377
389
|
// Construct a positive test set
|
378
390
|
auto queryVecs = faiss::gpu::randVecs(opt.numQuery, opt.dim);
|
@@ -419,7 +431,7 @@ TEST(TestGpuIndexIVFFlat, QueryNaN) {
|
|
419
431
|
|
420
432
|
faiss::gpu::GpuIndexIVFFlat gpuIndex(
|
421
433
|
&res, opt.dim, opt.numCentroids, faiss::METRIC_L2, config);
|
422
|
-
gpuIndex.
|
434
|
+
gpuIndex.nprobe = opt.nprobe;
|
423
435
|
|
424
436
|
gpuIndex.train(opt.numTrain, trainVecs.data());
|
425
437
|
gpuIndex.add(opt.numAdd, addVecs.data());
|
@@ -429,7 +441,7 @@ TEST(TestGpuIndexIVFFlat, QueryNaN) {
|
|
429
441
|
numQuery * opt.dim, std::numeric_limits<float>::quiet_NaN());
|
430
442
|
|
431
443
|
std::vector<float> distances(numQuery * opt.k, 0);
|
432
|
-
std::vector<faiss::
|
444
|
+
std::vector<faiss::idx_t> indices(numQuery * opt.k, 0);
|
433
445
|
|
434
446
|
gpuIndex.search(
|
435
447
|
numQuery, nans.data(), opt.k, distances.data(), indices.data());
|
@@ -457,7 +469,7 @@ TEST(TestGpuIndexIVFFlat, AddNaN) {
|
|
457
469
|
|
458
470
|
faiss::gpu::GpuIndexIVFFlat gpuIndex(
|
459
471
|
&res, opt.dim, opt.numCentroids, faiss::METRIC_L2, config);
|
460
|
-
gpuIndex.
|
472
|
+
gpuIndex.nprobe = opt.nprobe;
|
461
473
|
|
462
474
|
int numNans = 10;
|
463
475
|
std::vector<float> nans(
|
@@ -478,7 +490,7 @@ TEST(TestGpuIndexIVFFlat, AddNaN) {
|
|
478
490
|
|
479
491
|
std::vector<float> queryVecs = faiss::gpu::randVecs(opt.numQuery, opt.dim);
|
480
492
|
std::vector<float> distance(opt.numQuery * opt.k, 0);
|
481
|
-
std::vector<faiss::
|
493
|
+
std::vector<faiss::idx_t> indices(opt.numQuery * opt.k, 0);
|
482
494
|
|
483
495
|
// should not crash
|
484
496
|
gpuIndex.search(
|
@@ -531,7 +543,65 @@ TEST(TestGpuIndexIVFFlat, UnifiedMemory) {
|
|
531
543
|
faiss::gpu::GpuIndexIVFFlat gpuIndex(
|
532
544
|
&res, dim, numCentroids, faiss::METRIC_L2, config);
|
533
545
|
gpuIndex.copyFrom(&cpuIndex);
|
534
|
-
gpuIndex.
|
546
|
+
gpuIndex.nprobe = nprobe;
|
547
|
+
|
548
|
+
faiss::gpu::compareIndices(
|
549
|
+
cpuIndex,
|
550
|
+
gpuIndex,
|
551
|
+
numQuery,
|
552
|
+
dim,
|
553
|
+
k,
|
554
|
+
"Unified Memory",
|
555
|
+
kF32MaxRelErr,
|
556
|
+
0.1f,
|
557
|
+
0.015f);
|
558
|
+
}
|
559
|
+
|
560
|
+
TEST(TestGpuIndexIVFFlat, LongIVFList) {
|
561
|
+
int device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
|
562
|
+
|
563
|
+
// Skip this device if we do not have sufficient memory
|
564
|
+
constexpr size_t kMem = size_t(24) * 1024 * 1024 * 1024;
|
565
|
+
|
566
|
+
if (faiss::gpu::getFreeMemory(device) < kMem) {
|
567
|
+
std::cout << "TestGpuIndexIVFFlat.LongIVFList: skipping due "
|
568
|
+
"to insufficient device memory\n";
|
569
|
+
return;
|
570
|
+
}
|
571
|
+
|
572
|
+
std::cout << "Running LongIVFList test\n";
|
573
|
+
|
574
|
+
// Test functionality where a single IVF list has more than 2B code values
|
575
|
+
int dim = 64;
|
576
|
+
|
577
|
+
int numCentroids = 1;
|
578
|
+
size_t numAdd = (size_t(1024) * 1024 * 1024 * 2 + 100000) / dim;
|
579
|
+
size_t numTrain = 100;
|
580
|
+
int numQuery = 5;
|
581
|
+
int k = 10;
|
582
|
+
|
583
|
+
std::vector<float> trainVecs = faiss::gpu::randVecs(numTrain, dim);
|
584
|
+
std::vector<float> addVecs = faiss::gpu::randVecs(numAdd, dim);
|
585
|
+
|
586
|
+
faiss::IndexFlatL2 quantizer(dim);
|
587
|
+
faiss::IndexIVFFlat cpuIndex(
|
588
|
+
&quantizer, dim, numCentroids, faiss::METRIC_L2);
|
589
|
+
|
590
|
+
cpuIndex.train(numTrain, trainVecs.data());
|
591
|
+
cpuIndex.add(numAdd, addVecs.data());
|
592
|
+
cpuIndex.nprobe = 1;
|
593
|
+
|
594
|
+
faiss::gpu::StandardGpuResources res;
|
595
|
+
res.noTempMemory();
|
596
|
+
|
597
|
+
faiss::gpu::GpuIndexIVFFlatConfig config;
|
598
|
+
config.device = device;
|
599
|
+
|
600
|
+
faiss::gpu::GpuIndexIVFFlat gpuIndex(
|
601
|
+
&res, dim, numCentroids, faiss::METRIC_L2, config);
|
602
|
+
gpuIndex.train(numTrain, trainVecs.data());
|
603
|
+
gpuIndex.add(numAdd, addVecs.data());
|
604
|
+
gpuIndex.nprobe = 1;
|
535
605
|
|
536
606
|
faiss::gpu::compareIndices(
|
537
607
|
cpuIndex,
|
@@ -137,7 +137,58 @@ TEST(TestGpuIndexIVFPQ, Query_L2) {
|
|
137
137
|
config.useFloat16LookupTables = opt.useFloat16;
|
138
138
|
|
139
139
|
faiss::gpu::GpuIndexIVFPQ gpuIndex(&res, &cpuIndex, config);
|
140
|
-
gpuIndex.
|
140
|
+
gpuIndex.nprobe = opt.nprobe;
|
141
|
+
|
142
|
+
faiss::gpu::compareIndices(
|
143
|
+
cpuIndex,
|
144
|
+
gpuIndex,
|
145
|
+
opt.numQuery,
|
146
|
+
opt.dim,
|
147
|
+
opt.k,
|
148
|
+
opt.toString(),
|
149
|
+
opt.getCompareEpsilon(),
|
150
|
+
opt.getPctMaxDiff1(),
|
151
|
+
opt.getPctMaxDiffN());
|
152
|
+
}
|
153
|
+
}
|
154
|
+
|
155
|
+
// Large batch sizes (>= 65536) should also work
|
156
|
+
TEST(TestGpuIndexIVFPQ, LargeBatch) {
|
157
|
+
for (bool usePrecomputed : {false, true}) {
|
158
|
+
Options opt;
|
159
|
+
|
160
|
+
// override for large sizes
|
161
|
+
opt.dim = 4;
|
162
|
+
opt.numQuery = 100000;
|
163
|
+
opt.codes = 2;
|
164
|
+
|
165
|
+
std::vector<float> trainVecs =
|
166
|
+
faiss::gpu::randVecs(opt.numTrain, opt.dim);
|
167
|
+
std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
|
168
|
+
|
169
|
+
faiss::IndexFlatL2 coarseQuantizer(opt.dim);
|
170
|
+
faiss::IndexIVFPQ cpuIndex(
|
171
|
+
&coarseQuantizer,
|
172
|
+
opt.dim,
|
173
|
+
opt.numCentroids,
|
174
|
+
opt.codes,
|
175
|
+
opt.bitsPerCode);
|
176
|
+
cpuIndex.nprobe = opt.nprobe;
|
177
|
+
cpuIndex.train(opt.numTrain, trainVecs.data());
|
178
|
+
cpuIndex.add(opt.numAdd, addVecs.data());
|
179
|
+
|
180
|
+
// Use the default temporary memory management to test the memory
|
181
|
+
// manager
|
182
|
+
faiss::gpu::StandardGpuResources res;
|
183
|
+
|
184
|
+
faiss::gpu::GpuIndexIVFPQConfig config;
|
185
|
+
config.device = opt.device;
|
186
|
+
config.usePrecomputedTables = usePrecomputed;
|
187
|
+
config.indicesOptions = opt.indicesOpt;
|
188
|
+
config.useFloat16LookupTables = false;
|
189
|
+
|
190
|
+
faiss::gpu::GpuIndexIVFPQ gpuIndex(&res, &cpuIndex, config);
|
191
|
+
gpuIndex.nprobe = opt.nprobe;
|
141
192
|
|
142
193
|
faiss::gpu::compareIndices(
|
143
194
|
cpuIndex,
|
@@ -189,7 +240,7 @@ void testMMCodeDistance(faiss::MetricType mt) {
|
|
189
240
|
config.flatConfig.useFloat16 = (tries % 2 == 1);
|
190
241
|
|
191
242
|
faiss::gpu::GpuIndexIVFPQ gpuIndex(&res, &cpuIndex, config);
|
192
|
-
gpuIndex.
|
243
|
+
gpuIndex.nprobe = opt.nprobe;
|
193
244
|
|
194
245
|
faiss::gpu::compareIndices(
|
195
246
|
cpuIndex,
|
@@ -238,7 +289,7 @@ void testMMCodeDistance(faiss::MetricType mt) {
|
|
238
289
|
config.useFloat16LookupTables = (dimPerSubQ == 7);
|
239
290
|
|
240
291
|
faiss::gpu::GpuIndexIVFPQ gpuIndex(&res, &cpuIndex, config);
|
241
|
-
gpuIndex.
|
292
|
+
gpuIndex.nprobe = opt.nprobe;
|
242
293
|
|
243
294
|
faiss::gpu::compareIndices(
|
244
295
|
cpuIndex,
|
@@ -293,7 +344,7 @@ TEST(TestGpuIndexIVFPQ, Query_IP) {
|
|
293
344
|
config.useFloat16LookupTables = opt.useFloat16;
|
294
345
|
|
295
346
|
faiss::gpu::GpuIndexIVFPQ gpuIndex(&res, &cpuIndex, config);
|
296
|
-
gpuIndex.
|
347
|
+
gpuIndex.nprobe = opt.nprobe;
|
297
348
|
|
298
349
|
faiss::gpu::compareIndices(
|
299
350
|
cpuIndex,
|
@@ -335,7 +386,7 @@ TEST(TestGpuIndexIVFPQ, Float16Coarse) {
|
|
335
386
|
config.useFloat16LookupTables = opt.useFloat16;
|
336
387
|
|
337
388
|
faiss::gpu::GpuIndexIVFPQ gpuIndex(&res, &cpuIndex, config);
|
338
|
-
gpuIndex.
|
389
|
+
gpuIndex.nprobe = opt.nprobe;
|
339
390
|
|
340
391
|
gpuIndex.add(opt.numAdd, addVecs.data());
|
341
392
|
cpuIndex.add(opt.numAdd, addVecs.data());
|
@@ -381,7 +432,7 @@ TEST(TestGpuIndexIVFPQ, Add_L2) {
|
|
381
432
|
config.useFloat16LookupTables = opt.useFloat16;
|
382
433
|
|
383
434
|
faiss::gpu::GpuIndexIVFPQ gpuIndex(&res, &cpuIndex, config);
|
384
|
-
gpuIndex.
|
435
|
+
gpuIndex.nprobe = opt.nprobe;
|
385
436
|
|
386
437
|
gpuIndex.add(opt.numAdd, addVecs.data());
|
387
438
|
cpuIndex.add(opt.numAdd, addVecs.data());
|
@@ -429,7 +480,7 @@ TEST(TestGpuIndexIVFPQ, Add_IP) {
|
|
429
480
|
config.useFloat16LookupTables = opt.useFloat16;
|
430
481
|
|
431
482
|
faiss::gpu::GpuIndexIVFPQ gpuIndex(&res, &cpuIndex, config);
|
432
|
-
gpuIndex.
|
483
|
+
gpuIndex.nprobe = opt.nprobe;
|
433
484
|
|
434
485
|
gpuIndex.add(opt.numAdd, addVecs.data());
|
435
486
|
cpuIndex.add(opt.numAdd, addVecs.data());
|
@@ -472,7 +523,7 @@ TEST(TestGpuIndexIVFPQ, CopyTo) {
|
|
472
523
|
opt.bitsPerCode,
|
473
524
|
faiss::METRIC_L2,
|
474
525
|
config);
|
475
|
-
gpuIndex.
|
526
|
+
gpuIndex.nprobe = opt.nprobe;
|
476
527
|
gpuIndex.train(opt.numTrain, trainVecs.data());
|
477
528
|
gpuIndex.add(opt.numAdd, addVecs.data());
|
478
529
|
|
@@ -488,7 +539,7 @@ TEST(TestGpuIndexIVFPQ, CopyTo) {
|
|
488
539
|
EXPECT_EQ(cpuIndex.d, gpuIndex.d);
|
489
540
|
EXPECT_EQ(cpuIndex.d, opt.dim);
|
490
541
|
EXPECT_EQ(cpuIndex.nlist, gpuIndex.getNumLists());
|
491
|
-
EXPECT_EQ(cpuIndex.nprobe, gpuIndex.
|
542
|
+
EXPECT_EQ(cpuIndex.nprobe, gpuIndex.nprobe);
|
492
543
|
EXPECT_EQ(cpuIndex.pq.M, gpuIndex.getNumSubQuantizers());
|
493
544
|
EXPECT_EQ(gpuIndex.getNumSubQuantizers(), opt.codes);
|
494
545
|
EXPECT_EQ(cpuIndex.pq.nbits, gpuIndex.getBitsPerCode());
|
@@ -538,7 +589,7 @@ TEST(TestGpuIndexIVFPQ, CopyFrom) {
|
|
538
589
|
// Use garbage values to see if we overwrite them
|
539
590
|
faiss::gpu::GpuIndexIVFPQ gpuIndex(
|
540
591
|
&res, 1, 1, 1, 8, faiss::METRIC_L2, config);
|
541
|
-
gpuIndex.
|
592
|
+
gpuIndex.nprobe = 1;
|
542
593
|
|
543
594
|
gpuIndex.copyFrom(&cpuIndex);
|
544
595
|
|
@@ -549,7 +600,7 @@ TEST(TestGpuIndexIVFPQ, CopyFrom) {
|
|
549
600
|
EXPECT_EQ(cpuIndex.d, gpuIndex.d);
|
550
601
|
EXPECT_EQ(cpuIndex.d, opt.dim);
|
551
602
|
EXPECT_EQ(cpuIndex.nlist, gpuIndex.getNumLists());
|
552
|
-
EXPECT_EQ(cpuIndex.nprobe, gpuIndex.
|
603
|
+
EXPECT_EQ(cpuIndex.nprobe, gpuIndex.nprobe);
|
553
604
|
EXPECT_EQ(cpuIndex.pq.M, gpuIndex.getNumSubQuantizers());
|
554
605
|
EXPECT_EQ(gpuIndex.getNumSubQuantizers(), opt.codes);
|
555
606
|
EXPECT_EQ(cpuIndex.pq.nbits, gpuIndex.getBitsPerCode());
|
@@ -594,7 +645,7 @@ TEST(TestGpuIndexIVFPQ, QueryNaN) {
|
|
594
645
|
faiss::METRIC_L2,
|
595
646
|
config);
|
596
647
|
|
597
|
-
gpuIndex.
|
648
|
+
gpuIndex.nprobe = opt.nprobe;
|
598
649
|
|
599
650
|
gpuIndex.train(opt.numTrain, trainVecs.data());
|
600
651
|
gpuIndex.add(opt.numAdd, addVecs.data());
|
@@ -604,7 +655,7 @@ TEST(TestGpuIndexIVFPQ, QueryNaN) {
|
|
604
655
|
numQuery * opt.dim, std::numeric_limits<float>::quiet_NaN());
|
605
656
|
|
606
657
|
std::vector<float> distances(numQuery * opt.k, 0);
|
607
|
-
std::vector<faiss::
|
658
|
+
std::vector<faiss::idx_t> indices(numQuery * opt.k, 0);
|
608
659
|
|
609
660
|
gpuIndex.search(
|
610
661
|
numQuery, nans.data(), opt.k, distances.data(), indices.data());
|
@@ -640,7 +691,7 @@ TEST(TestGpuIndexIVFPQ, AddNaN) {
|
|
640
691
|
faiss::METRIC_L2,
|
641
692
|
config);
|
642
693
|
|
643
|
-
gpuIndex.
|
694
|
+
gpuIndex.nprobe = opt.nprobe;
|
644
695
|
|
645
696
|
int numNans = 10;
|
646
697
|
std::vector<float> nans(
|
@@ -660,7 +711,7 @@ TEST(TestGpuIndexIVFPQ, AddNaN) {
|
|
660
711
|
|
661
712
|
std::vector<float> queryVecs = faiss::gpu::randVecs(opt.numQuery, opt.dim);
|
662
713
|
std::vector<float> distance(opt.numQuery * opt.k, 0);
|
663
|
-
std::vector<faiss::
|
714
|
+
std::vector<faiss::idx_t> indices(opt.numQuery * opt.k, 0);
|
664
715
|
|
665
716
|
// should not crash
|
666
717
|
gpuIndex.search(
|
@@ -721,7 +772,7 @@ TEST(TestGpuIndexIVFPQ, UnifiedMemory) {
|
|
721
772
|
faiss::METRIC_L2,
|
722
773
|
config);
|
723
774
|
gpuIndex.copyFrom(&cpuIndex);
|
724
|
-
gpuIndex.
|
775
|
+
gpuIndex.nprobe = nprobe;
|
725
776
|
|
726
777
|
faiss::gpu::compareIndices(
|
727
778
|
cpuIndex,
|
@@ -79,7 +79,7 @@ void runCopyToTest(faiss::ScalarQuantizer::QuantizerType qtype) {
|
|
79
79
|
&res, opt.dim, opt.numCentroids, qtype, METRIC_L2, true, config);
|
80
80
|
gpuIndex.train(opt.numTrain, trainVecs.data());
|
81
81
|
gpuIndex.add(opt.numAdd, addVecs.data());
|
82
|
-
gpuIndex.
|
82
|
+
gpuIndex.nprobe = opt.nprobe;
|
83
83
|
|
84
84
|
// use garbage values to see if we overwrite then
|
85
85
|
IndexFlatL2 cpuQuantizer(1);
|
@@ -100,7 +100,7 @@ void runCopyToTest(faiss::ScalarQuantizer::QuantizerType qtype) {
|
|
100
100
|
EXPECT_EQ(cpuIndex.quantizer->d, gpuIndex.quantizer->d);
|
101
101
|
EXPECT_EQ(cpuIndex.d, opt.dim);
|
102
102
|
EXPECT_EQ(cpuIndex.nlist, gpuIndex.getNumLists());
|
103
|
-
EXPECT_EQ(cpuIndex.nprobe, gpuIndex.
|
103
|
+
EXPECT_EQ(cpuIndex.nprobe, gpuIndex.nprobe);
|
104
104
|
|
105
105
|
testIVFEquality(cpuIndex, gpuIndex);
|
106
106
|
|
@@ -172,7 +172,7 @@ void runCopyFromTest(faiss::ScalarQuantizer::QuantizerType qtype) {
|
|
172
172
|
METRIC_L2,
|
173
173
|
false,
|
174
174
|
config);
|
175
|
-
gpuIndex.
|
175
|
+
gpuIndex.nprobe = 1;
|
176
176
|
|
177
177
|
gpuIndex.copyFrom(&cpuIndex);
|
178
178
|
|
@@ -182,7 +182,7 @@ void runCopyFromTest(faiss::ScalarQuantizer::QuantizerType qtype) {
|
|
182
182
|
EXPECT_EQ(cpuIndex.d, gpuIndex.d);
|
183
183
|
EXPECT_EQ(cpuIndex.d, opt.dim);
|
184
184
|
EXPECT_EQ(cpuIndex.nlist, gpuIndex.getNumLists());
|
185
|
-
EXPECT_EQ(cpuIndex.nprobe, gpuIndex.
|
185
|
+
EXPECT_EQ(cpuIndex.nprobe, gpuIndex.nprobe);
|
186
186
|
|
187
187
|
testIVFEquality(cpuIndex, gpuIndex);
|
188
188
|
|
@@ -97,7 +97,7 @@ void compareIndices(
|
|
97
97
|
float pctMaxDiffN) {
|
98
98
|
// Compare
|
99
99
|
std::vector<float> refDistance(numQuery * k, 0);
|
100
|
-
std::vector<faiss::
|
100
|
+
std::vector<faiss::idx_t> refIndices(numQuery * k, -1);
|
101
101
|
refIndex.search(
|
102
102
|
numQuery,
|
103
103
|
queryVecs.data(),
|
@@ -106,7 +106,7 @@ void compareIndices(
|
|
106
106
|
refIndices.data());
|
107
107
|
|
108
108
|
std::vector<float> testDistance(numQuery * k, 0);
|
109
|
-
std::vector<faiss::
|
109
|
+
std::vector<faiss::idx_t> testIndices(numQuery * k, -1);
|
110
110
|
testIndex.search(
|
111
111
|
numQuery,
|
112
112
|
queryVecs.data(),
|
@@ -162,9 +162,9 @@ inline T lookup(const T* p, int i, int j, int /*dim1*/, int dim2) {
|
|
162
162
|
|
163
163
|
void compareLists(
|
164
164
|
const float* refDist,
|
165
|
-
const faiss::
|
165
|
+
const faiss::idx_t* refInd,
|
166
166
|
const float* testDist,
|
167
|
-
const faiss::
|
167
|
+
const faiss::idx_t* testInd,
|
168
168
|
int dim1,
|
169
169
|
int dim2,
|
170
170
|
const std::string& configMsg,
|
@@ -181,10 +181,10 @@ void compareLists(
|
|
181
181
|
int numResults = dim1 * dim2;
|
182
182
|
|
183
183
|
// query -> {index -> result position}
|
184
|
-
std::vector<std::unordered_map<faiss::
|
184
|
+
std::vector<std::unordered_map<faiss::idx_t, int>> refIndexMap;
|
185
185
|
|
186
186
|
for (int query = 0; query < dim1; ++query) {
|
187
|
-
std::unordered_map<faiss::
|
187
|
+
std::unordered_map<faiss::idx_t, int> indices;
|
188
188
|
|
189
189
|
for (int result = 0; result < dim2; ++result) {
|
190
190
|
indices[lookup(refInd, query, result, dim1, dim2)] = result;
|
@@ -208,7 +208,7 @@ void compareLists(
|
|
208
208
|
|
209
209
|
for (int query = 0; query < dim1; ++query) {
|
210
210
|
std::vector<int> diffs;
|
211
|
-
std::set<faiss::
|
211
|
+
std::set<faiss::idx_t> uniqueIndices;
|
212
212
|
|
213
213
|
auto& indices = refIndexMap[query];
|
214
214
|
|
@@ -93,9 +93,9 @@ void compareIndices(
|
|
93
93
|
/// Display specific differences in the two (distance, index) lists
|
94
94
|
void compareLists(
|
95
95
|
const float* refDist,
|
96
|
-
const faiss::
|
96
|
+
const faiss::idx_t* refInd,
|
97
97
|
const float* testDist,
|
98
|
-
const faiss::
|
98
|
+
const faiss::idx_t* testInd,
|
99
99
|
int dim1,
|
100
100
|
int dim2,
|
101
101
|
const std::string& configMsg,
|
@@ -130,13 +130,13 @@ void testIVFEquality(A& cpuIndex, B& gpuIndex) {
|
|
130
130
|
EXPECT_EQ(cpuCodes, gpuCodes);
|
131
131
|
|
132
132
|
// Index equality
|
133
|
-
std::vector<
|
133
|
+
std::vector<idx_t> cpuIndices(cpuLists->list_size(i));
|
134
134
|
|
135
135
|
auto si = faiss::InvertedLists::ScopedIds(cpuLists, i);
|
136
136
|
std::memcpy(
|
137
137
|
cpuIndices.data(),
|
138
138
|
si.get(),
|
139
|
-
cpuLists->list_size(i) * sizeof(faiss::
|
139
|
+
cpuLists->list_size(i) * sizeof(faiss::idx_t));
|
140
140
|
EXPECT_EQ(cpuIndices, gpuIndex.getListIndices(i));
|
141
141
|
}
|
142
142
|
}
|
@@ -47,6 +47,12 @@ int getMaxThreads(int device);
|
|
47
47
|
/// Equivalent to getMaxThreads(getCurrentDevice())
|
48
48
|
int getMaxThreadsCurrentDevice();
|
49
49
|
|
50
|
+
/// Returns the maximum grid size for the given GPU device
|
51
|
+
dim3 getMaxGrid(int device);
|
52
|
+
|
53
|
+
/// Equivalent to getMaxGrid(getCurrentDevice())
|
54
|
+
dim3 getMaxGridCurrentDevice();
|
55
|
+
|
50
56
|
/// Returns the maximum smem available for the given GPU device
|
51
57
|
size_t getMaxSharedMemPerBlock(int device);
|
52
58
|
|
@@ -54,14 +54,7 @@ AdditiveQuantizer::AdditiveQuantizer(
|
|
54
54
|
: Quantizer(d),
|
55
55
|
M(nbits.size()),
|
56
56
|
nbits(nbits),
|
57
|
-
verbose(false),
|
58
|
-
is_trained(false),
|
59
|
-
max_mem_distances(5 * (size_t(1) << 30)), // 5 GiB
|
60
57
|
search_type(search_type) {
|
61
|
-
norm_max = norm_min = NAN;
|
62
|
-
tot_bits = 0;
|
63
|
-
total_codebook_size = 0;
|
64
|
-
only_8bit = false;
|
65
58
|
set_derived_values();
|
66
59
|
}
|
67
60
|
|
@@ -7,6 +7,7 @@
|
|
7
7
|
|
8
8
|
#pragma once
|
9
9
|
|
10
|
+
#include <cmath>
|
10
11
|
#include <cstdint>
|
11
12
|
#include <vector>
|
12
13
|
|
@@ -29,13 +30,13 @@ struct AdditiveQuantizer : Quantizer {
|
|
29
30
|
|
30
31
|
// derived values
|
31
32
|
std::vector<uint64_t> codebook_offsets;
|
32
|
-
size_t tot_bits; ///< total number of bits (indexes + norms)
|
33
|
-
size_t norm_bits; ///< bits allocated for the norms
|
34
|
-
size_t total_codebook_size; ///< size of the codebook in vectors
|
35
|
-
bool only_8bit;
|
33
|
+
size_t tot_bits = 0; ///< total number of bits (indexes + norms)
|
34
|
+
size_t norm_bits = 0; ///< bits allocated for the norms
|
35
|
+
size_t total_codebook_size = 0; ///< size of the codebook in vectors
|
36
|
+
bool only_8bit = false; ///< are all nbits = 8 (use faster decoder)
|
36
37
|
|
37
|
-
bool verbose; ///< verbose during training?
|
38
|
-
bool is_trained; ///< is trained or not
|
38
|
+
bool verbose = false; ///< verbose during training?
|
39
|
+
bool is_trained = false; ///< is trained or not
|
39
40
|
|
40
41
|
IndexFlat1D qnorm; ///< store and search norms
|
41
42
|
std::vector<float> norm_tabs; ///< store norms of codebook entries for 4-bit
|
@@ -43,7 +44,7 @@ struct AdditiveQuantizer : Quantizer {
|
|
43
44
|
|
44
45
|
/// norms and distance matrixes with beam search can get large, so use this
|
45
46
|
/// to control for the amount of memory that can be allocated
|
46
|
-
size_t max_mem_distances;
|
47
|
+
size_t max_mem_distances = 5 * (size_t(1) << 30);
|
47
48
|
|
48
49
|
/// encode a norm into norm_bits bits
|
49
50
|
uint64_t encode_norm(float norm) const;
|
@@ -145,7 +146,7 @@ struct AdditiveQuantizer : Quantizer {
|
|
145
146
|
Search_type_t search_type;
|
146
147
|
|
147
148
|
/// min/max for quantization of norms
|
148
|
-
float norm_min, norm_max;
|
149
|
+
float norm_min = NAN, norm_max = NAN;
|
149
150
|
|
150
151
|
template <bool is_IP, Search_type_t effective_search_type>
|
151
152
|
float compute_1_distance_LUT(const uint8_t* codes, const float* LUT) const;
|
@@ -157,7 +158,6 @@ struct AdditiveQuantizer : Quantizer {
|
|
157
158
|
* Support for exhaustive distance computations with all the centroids.
|
158
159
|
* Hence, the number of these centroids should not be too large.
|
159
160
|
****************************************************************************/
|
160
|
-
using idx_t = Index::idx_t;
|
161
161
|
|
162
162
|
/// decoding function for a code in a 64-bit word
|
163
163
|
void decode_64bit(idx_t n, float* x) const;
|
@@ -20,7 +20,7 @@ namespace faiss {
|
|
20
20
|
* RangeSearchResult
|
21
21
|
***********************************************************************/
|
22
22
|
|
23
|
-
RangeSearchResult::RangeSearchResult(
|
23
|
+
RangeSearchResult::RangeSearchResult(size_t nq, bool alloc_lims) : nq(nq) {
|
24
24
|
if (alloc_lims) {
|
25
25
|
lims = new size_t[nq + 1];
|
26
26
|
memset(lims, 0, sizeof(*lims) * (nq + 1));
|