faiss 0.3.0 → 0.3.2
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- checksums.yaml +4 -4
- data/CHANGELOG.md +9 -0
- data/LICENSE.txt +1 -1
- data/README.md +1 -1
- data/ext/faiss/extconf.rb +9 -2
- data/ext/faiss/index.cpp +1 -1
- data/ext/faiss/index_binary.cpp +2 -2
- data/ext/faiss/product_quantizer.cpp +1 -1
- data/lib/faiss/version.rb +1 -1
- data/vendor/faiss/faiss/AutoTune.cpp +7 -7
- data/vendor/faiss/faiss/AutoTune.h +1 -2
- data/vendor/faiss/faiss/Clustering.cpp +39 -22
- data/vendor/faiss/faiss/Clustering.h +40 -21
- data/vendor/faiss/faiss/IVFlib.cpp +26 -12
- data/vendor/faiss/faiss/Index.cpp +1 -1
- data/vendor/faiss/faiss/Index.h +40 -10
- data/vendor/faiss/faiss/Index2Layer.cpp +7 -7
- data/vendor/faiss/faiss/IndexAdditiveQuantizer.cpp +176 -166
- data/vendor/faiss/faiss/IndexAdditiveQuantizerFastScan.cpp +15 -15
- data/vendor/faiss/faiss/IndexBinary.cpp +9 -4
- data/vendor/faiss/faiss/IndexBinary.h +8 -19
- data/vendor/faiss/faiss/IndexBinaryFromFloat.cpp +2 -1
- data/vendor/faiss/faiss/IndexBinaryHNSW.cpp +24 -31
- data/vendor/faiss/faiss/IndexBinaryHNSW.h +1 -1
- data/vendor/faiss/faiss/IndexBinaryHash.cpp +25 -50
- data/vendor/faiss/faiss/IndexBinaryIVF.cpp +107 -188
- data/vendor/faiss/faiss/IndexFastScan.cpp +95 -146
- data/vendor/faiss/faiss/IndexFastScan.h +9 -8
- data/vendor/faiss/faiss/IndexFlat.cpp +206 -10
- data/vendor/faiss/faiss/IndexFlat.h +20 -1
- data/vendor/faiss/faiss/IndexFlatCodes.cpp +170 -5
- data/vendor/faiss/faiss/IndexFlatCodes.h +23 -4
- data/vendor/faiss/faiss/IndexHNSW.cpp +231 -382
- data/vendor/faiss/faiss/IndexHNSW.h +62 -49
- data/vendor/faiss/faiss/IndexIDMap.cpp +69 -28
- data/vendor/faiss/faiss/IndexIDMap.h +24 -2
- data/vendor/faiss/faiss/IndexIVF.cpp +162 -56
- data/vendor/faiss/faiss/IndexIVF.h +46 -6
- data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.cpp +33 -26
- data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.h +6 -2
- data/vendor/faiss/faiss/IndexIVFAdditiveQuantizerFastScan.cpp +19 -46
- data/vendor/faiss/faiss/IndexIVFAdditiveQuantizerFastScan.h +4 -3
- data/vendor/faiss/faiss/IndexIVFFastScan.cpp +502 -401
- data/vendor/faiss/faiss/IndexIVFFastScan.h +63 -26
- data/vendor/faiss/faiss/IndexIVFFlat.cpp +15 -5
- data/vendor/faiss/faiss/IndexIVFFlat.h +3 -2
- data/vendor/faiss/faiss/IndexIVFIndependentQuantizer.cpp +172 -0
- data/vendor/faiss/faiss/IndexIVFIndependentQuantizer.h +56 -0
- data/vendor/faiss/faiss/IndexIVFPQ.cpp +79 -125
- data/vendor/faiss/faiss/IndexIVFPQ.h +6 -7
- data/vendor/faiss/faiss/IndexIVFPQFastScan.cpp +39 -52
- data/vendor/faiss/faiss/IndexIVFPQFastScan.h +4 -3
- data/vendor/faiss/faiss/IndexIVFPQR.cpp +45 -29
- data/vendor/faiss/faiss/IndexIVFPQR.h +5 -2
- data/vendor/faiss/faiss/IndexIVFSpectralHash.cpp +25 -27
- data/vendor/faiss/faiss/IndexIVFSpectralHash.h +6 -6
- data/vendor/faiss/faiss/IndexLSH.cpp +14 -16
- data/vendor/faiss/faiss/IndexLattice.cpp +1 -19
- data/vendor/faiss/faiss/IndexLattice.h +3 -22
- data/vendor/faiss/faiss/IndexNNDescent.cpp +3 -33
- data/vendor/faiss/faiss/IndexNNDescent.h +1 -1
- data/vendor/faiss/faiss/IndexNSG.cpp +11 -27
- data/vendor/faiss/faiss/IndexNSG.h +11 -11
- data/vendor/faiss/faiss/IndexNeuralNetCodec.cpp +56 -0
- data/vendor/faiss/faiss/IndexNeuralNetCodec.h +49 -0
- data/vendor/faiss/faiss/IndexPQ.cpp +72 -88
- data/vendor/faiss/faiss/IndexPQ.h +1 -4
- data/vendor/faiss/faiss/IndexPQFastScan.cpp +1 -1
- data/vendor/faiss/faiss/IndexPreTransform.cpp +25 -31
- data/vendor/faiss/faiss/IndexPreTransform.h +1 -1
- data/vendor/faiss/faiss/IndexRefine.cpp +54 -24
- data/vendor/faiss/faiss/IndexRefine.h +7 -0
- data/vendor/faiss/faiss/IndexReplicas.cpp +23 -26
- data/vendor/faiss/faiss/IndexScalarQuantizer.cpp +25 -17
- data/vendor/faiss/faiss/IndexScalarQuantizer.h +6 -4
- data/vendor/faiss/faiss/IndexShards.cpp +21 -29
- data/vendor/faiss/faiss/IndexShardsIVF.cpp +1 -2
- data/vendor/faiss/faiss/MatrixStats.cpp +17 -32
- data/vendor/faiss/faiss/MatrixStats.h +21 -9
- data/vendor/faiss/faiss/MetaIndexes.cpp +35 -35
- data/vendor/faiss/faiss/MetricType.h +7 -2
- data/vendor/faiss/faiss/VectorTransform.cpp +13 -26
- data/vendor/faiss/faiss/VectorTransform.h +7 -7
- data/vendor/faiss/faiss/clone_index.cpp +15 -10
- data/vendor/faiss/faiss/clone_index.h +3 -0
- data/vendor/faiss/faiss/cppcontrib/detail/UintReader.h +95 -17
- data/vendor/faiss/faiss/cppcontrib/factory_tools.cpp +152 -0
- data/vendor/faiss/faiss/cppcontrib/factory_tools.h +24 -0
- data/vendor/faiss/faiss/cppcontrib/sa_decode/Level2-inl.h +83 -30
- data/vendor/faiss/faiss/gpu/GpuCloner.cpp +123 -8
- data/vendor/faiss/faiss/gpu/GpuCloner.h +22 -0
- data/vendor/faiss/faiss/gpu/GpuClonerOptions.h +13 -0
- data/vendor/faiss/faiss/gpu/GpuDistance.h +46 -38
- data/vendor/faiss/faiss/gpu/GpuFaissAssert.h +1 -1
- data/vendor/faiss/faiss/gpu/GpuIndex.h +30 -12
- data/vendor/faiss/faiss/gpu/GpuIndexCagra.h +282 -0
- data/vendor/faiss/faiss/gpu/GpuIndexFlat.h +4 -4
- data/vendor/faiss/faiss/gpu/GpuIndexIVF.h +14 -9
- data/vendor/faiss/faiss/gpu/GpuIndexIVFFlat.h +20 -3
- data/vendor/faiss/faiss/gpu/GpuIndexIVFPQ.h +22 -11
- data/vendor/faiss/faiss/gpu/GpuIndexIVFScalarQuantizer.h +1 -3
- data/vendor/faiss/faiss/gpu/GpuResources.cpp +24 -3
- data/vendor/faiss/faiss/gpu/GpuResources.h +39 -11
- data/vendor/faiss/faiss/gpu/StandardGpuResources.cpp +142 -17
- data/vendor/faiss/faiss/gpu/StandardGpuResources.h +57 -3
- data/vendor/faiss/faiss/gpu/impl/InterleavedCodes.cpp +26 -21
- data/vendor/faiss/faiss/gpu/perf/PerfClustering.cpp +7 -1
- data/vendor/faiss/faiss/gpu/test/TestCodePacking.cpp +8 -5
- data/vendor/faiss/faiss/gpu/test/TestGpuIndexBinaryFlat.cpp +25 -0
- data/vendor/faiss/faiss/gpu/test/TestGpuIndexFlat.cpp +129 -9
- data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFFlat.cpp +332 -40
- data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFPQ.cpp +299 -208
- data/vendor/faiss/faiss/gpu/test/TestGpuMemoryException.cpp +1 -0
- 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/gpu/utils/RaftUtils.h +75 -0
- data/vendor/faiss/faiss/gpu/utils/Timer.cpp +4 -1
- data/vendor/faiss/faiss/gpu/utils/Timer.h +1 -1
- data/vendor/faiss/faiss/impl/AdditiveQuantizer.cpp +3 -1
- data/vendor/faiss/faiss/impl/AdditiveQuantizer.h +5 -5
- data/vendor/faiss/faiss/impl/AuxIndexStructures.cpp +26 -1
- data/vendor/faiss/faiss/impl/AuxIndexStructures.h +10 -3
- data/vendor/faiss/faiss/impl/DistanceComputer.h +70 -1
- data/vendor/faiss/faiss/impl/FaissAssert.h +4 -2
- data/vendor/faiss/faiss/impl/FaissException.h +13 -34
- data/vendor/faiss/faiss/impl/HNSW.cpp +605 -186
- data/vendor/faiss/faiss/impl/HNSW.h +52 -30
- data/vendor/faiss/faiss/impl/IDSelector.h +4 -4
- data/vendor/faiss/faiss/impl/LocalSearchQuantizer.cpp +11 -9
- data/vendor/faiss/faiss/impl/LookupTableScaler.h +34 -0
- data/vendor/faiss/faiss/impl/NNDescent.cpp +42 -27
- data/vendor/faiss/faiss/impl/NSG.cpp +0 -29
- data/vendor/faiss/faiss/impl/NSG.h +1 -1
- data/vendor/faiss/faiss/impl/PolysemousTraining.cpp +14 -12
- data/vendor/faiss/faiss/impl/ProductAdditiveQuantizer.h +1 -1
- data/vendor/faiss/faiss/impl/ProductQuantizer.cpp +25 -22
- data/vendor/faiss/faiss/impl/ProductQuantizer.h +6 -2
- data/vendor/faiss/faiss/impl/Quantizer.h +1 -1
- data/vendor/faiss/faiss/impl/ResidualQuantizer.cpp +27 -1015
- data/vendor/faiss/faiss/impl/ResidualQuantizer.h +5 -63
- data/vendor/faiss/faiss/impl/ResultHandler.h +347 -172
- data/vendor/faiss/faiss/impl/ScalarQuantizer.cpp +1104 -147
- data/vendor/faiss/faiss/impl/ScalarQuantizer.h +3 -8
- data/vendor/faiss/faiss/impl/code_distance/code_distance-avx2.h +285 -42
- data/vendor/faiss/faiss/impl/code_distance/code_distance-avx512.h +248 -0
- data/vendor/faiss/faiss/impl/code_distance/code_distance-generic.h +21 -14
- data/vendor/faiss/faiss/impl/code_distance/code_distance.h +22 -12
- data/vendor/faiss/faiss/impl/index_read.cpp +74 -34
- data/vendor/faiss/faiss/impl/index_read_utils.h +37 -0
- data/vendor/faiss/faiss/impl/index_write.cpp +88 -51
- data/vendor/faiss/faiss/impl/io.cpp +23 -15
- data/vendor/faiss/faiss/impl/io.h +4 -4
- data/vendor/faiss/faiss/impl/io_macros.h +6 -0
- data/vendor/faiss/faiss/impl/lattice_Zn.cpp +1 -1
- data/vendor/faiss/faiss/impl/platform_macros.h +40 -1
- data/vendor/faiss/faiss/impl/pq4_fast_scan.cpp +14 -0
- data/vendor/faiss/faiss/impl/pq4_fast_scan.h +7 -6
- data/vendor/faiss/faiss/impl/pq4_fast_scan_search_1.cpp +52 -38
- data/vendor/faiss/faiss/impl/pq4_fast_scan_search_qbs.cpp +487 -49
- data/vendor/faiss/faiss/impl/residual_quantizer_encode_steps.cpp +960 -0
- data/vendor/faiss/faiss/impl/residual_quantizer_encode_steps.h +176 -0
- data/vendor/faiss/faiss/impl/simd_result_handlers.h +481 -225
- data/vendor/faiss/faiss/index_factory.cpp +41 -20
- data/vendor/faiss/faiss/index_io.h +12 -5
- data/vendor/faiss/faiss/invlists/BlockInvertedLists.cpp +28 -8
- data/vendor/faiss/faiss/invlists/BlockInvertedLists.h +3 -0
- data/vendor/faiss/faiss/invlists/DirectMap.cpp +10 -2
- data/vendor/faiss/faiss/invlists/InvertedLists.cpp +73 -17
- data/vendor/faiss/faiss/invlists/InvertedLists.h +26 -8
- data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.cpp +24 -9
- data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.h +2 -1
- data/vendor/faiss/faiss/python/python_callbacks.cpp +4 -4
- data/vendor/faiss/faiss/utils/Heap.cpp +3 -1
- data/vendor/faiss/faiss/utils/Heap.h +105 -0
- data/vendor/faiss/faiss/utils/NeuralNet.cpp +342 -0
- data/vendor/faiss/faiss/utils/NeuralNet.h +147 -0
- data/vendor/faiss/faiss/utils/WorkerThread.h +1 -0
- data/vendor/faiss/faiss/utils/bf16.h +36 -0
- data/vendor/faiss/faiss/utils/distances.cpp +147 -123
- data/vendor/faiss/faiss/utils/distances.h +86 -9
- data/vendor/faiss/faiss/utils/distances_fused/avx512.cpp +5 -5
- data/vendor/faiss/faiss/utils/distances_fused/avx512.h +2 -2
- data/vendor/faiss/faiss/utils/distances_fused/distances_fused.cpp +2 -2
- data/vendor/faiss/faiss/utils/distances_fused/distances_fused.h +1 -1
- data/vendor/faiss/faiss/utils/distances_fused/simdlib_based.cpp +5 -5
- data/vendor/faiss/faiss/utils/distances_fused/simdlib_based.h +1 -1
- data/vendor/faiss/faiss/utils/distances_simd.cpp +1589 -243
- data/vendor/faiss/faiss/utils/extra_distances-inl.h +70 -0
- data/vendor/faiss/faiss/utils/extra_distances.cpp +85 -137
- data/vendor/faiss/faiss/utils/extra_distances.h +3 -2
- data/vendor/faiss/faiss/utils/fp16-arm.h +29 -0
- data/vendor/faiss/faiss/utils/fp16.h +2 -0
- data/vendor/faiss/faiss/utils/hamming.cpp +163 -111
- data/vendor/faiss/faiss/utils/hamming.h +58 -0
- data/vendor/faiss/faiss/utils/hamming_distance/avx2-inl.h +16 -89
- data/vendor/faiss/faiss/utils/hamming_distance/common.h +1 -0
- data/vendor/faiss/faiss/utils/hamming_distance/generic-inl.h +19 -88
- data/vendor/faiss/faiss/utils/hamming_distance/hamdis-inl.h +58 -0
- data/vendor/faiss/faiss/utils/hamming_distance/neon-inl.h +14 -104
- data/vendor/faiss/faiss/utils/partitioning.cpp +3 -4
- data/vendor/faiss/faiss/utils/prefetch.h +77 -0
- data/vendor/faiss/faiss/utils/quantize_lut.cpp +0 -14
- data/vendor/faiss/faiss/utils/random.cpp +43 -0
- data/vendor/faiss/faiss/utils/random.h +25 -0
- data/vendor/faiss/faiss/utils/simdlib.h +10 -1
- data/vendor/faiss/faiss/utils/simdlib_avx2.h +0 -6
- data/vendor/faiss/faiss/utils/simdlib_avx512.h +296 -0
- data/vendor/faiss/faiss/utils/simdlib_neon.h +77 -79
- data/vendor/faiss/faiss/utils/simdlib_ppc64.h +1084 -0
- data/vendor/faiss/faiss/utils/sorting.cpp +140 -5
- data/vendor/faiss/faiss/utils/sorting.h +27 -0
- data/vendor/faiss/faiss/utils/transpose/transpose-avx512-inl.h +176 -0
- data/vendor/faiss/faiss/utils/utils.cpp +120 -7
- data/vendor/faiss/faiss/utils/utils.h +60 -20
- metadata +23 -4
- data/vendor/faiss/faiss/impl/code_distance/code_distance_avx512.h +0 -102
|
@@ -4,10 +4,26 @@
|
|
|
4
4
|
* This source code is licensed under the MIT license found in the
|
|
5
5
|
* LICENSE file in the root directory of this source tree.
|
|
6
6
|
*/
|
|
7
|
+
/*
|
|
8
|
+
* Copyright (c) 2023, NVIDIA CORPORATION.
|
|
9
|
+
*
|
|
10
|
+
* Licensed under the Apache License, Version 2.0 (the "License");
|
|
11
|
+
* you may not use this file except in compliance with the License.
|
|
12
|
+
* You may obtain a copy of the License at
|
|
13
|
+
*
|
|
14
|
+
* http://www.apache.org/licenses/LICENSE-2.0
|
|
15
|
+
*
|
|
16
|
+
* Unless required by applicable law or agreed to in writing, software
|
|
17
|
+
* distributed under the License is distributed on an "AS IS" BASIS,
|
|
18
|
+
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
19
|
+
* See the License for the specific language governing permissions and
|
|
20
|
+
* limitations under the License.
|
|
21
|
+
*/
|
|
7
22
|
|
|
8
23
|
#include <faiss/IndexFlat.h>
|
|
9
24
|
#include <faiss/IndexIVFFlat.h>
|
|
10
25
|
#include <faiss/gpu/GpuIndexIVFFlat.h>
|
|
26
|
+
#include <faiss/gpu/GpuIndicesOptions.h>
|
|
11
27
|
#include <faiss/gpu/StandardGpuResources.h>
|
|
12
28
|
#include <faiss/gpu/test/TestUtils.h>
|
|
13
29
|
#include <faiss/gpu/utils/DeviceUtils.h>
|
|
@@ -40,6 +56,8 @@ struct Options {
|
|
|
40
56
|
faiss::gpu::INDICES_64_BIT});
|
|
41
57
|
|
|
42
58
|
device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
|
|
59
|
+
|
|
60
|
+
useRaft = false;
|
|
43
61
|
}
|
|
44
62
|
|
|
45
63
|
std::string toString() const {
|
|
@@ -47,7 +65,7 @@ struct Options {
|
|
|
47
65
|
str << "IVFFlat device " << device << " numVecs " << numAdd << " dim "
|
|
48
66
|
<< dim << " numCentroids " << numCentroids << " nprobe " << nprobe
|
|
49
67
|
<< " numQuery " << numQuery << " k " << k << " indicesOpt "
|
|
50
|
-
<< indicesOpt;
|
|
68
|
+
<< indicesOpt << " useRaft " << useRaft;
|
|
51
69
|
|
|
52
70
|
return str.str();
|
|
53
71
|
}
|
|
@@ -61,6 +79,7 @@ struct Options {
|
|
|
61
79
|
int k;
|
|
62
80
|
int device;
|
|
63
81
|
faiss::gpu::IndicesOptions indicesOpt;
|
|
82
|
+
bool useRaft;
|
|
64
83
|
};
|
|
65
84
|
|
|
66
85
|
void queryTest(
|
|
@@ -91,6 +110,7 @@ void queryTest(
|
|
|
91
110
|
config.device = opt.device;
|
|
92
111
|
config.indicesOptions = opt.indicesOpt;
|
|
93
112
|
config.flatConfig.useFloat16 = useFloat16CoarseQuantizer;
|
|
113
|
+
config.use_raft = opt.useRaft;
|
|
94
114
|
|
|
95
115
|
faiss::gpu::GpuIndexIVFFlat gpuIndex(
|
|
96
116
|
&res, cpuIndex.d, cpuIndex.nlist, cpuIndex.metric_type, config);
|
|
@@ -114,7 +134,10 @@ void queryTest(
|
|
|
114
134
|
}
|
|
115
135
|
}
|
|
116
136
|
|
|
117
|
-
void addTest(
|
|
137
|
+
void addTest(
|
|
138
|
+
faiss::MetricType metricType,
|
|
139
|
+
bool useFloat16CoarseQuantizer,
|
|
140
|
+
bool useRaft) {
|
|
118
141
|
for (int tries = 0; tries < 2; ++tries) {
|
|
119
142
|
Options opt;
|
|
120
143
|
|
|
@@ -138,8 +161,10 @@ void addTest(faiss::MetricType metricType, bool useFloat16CoarseQuantizer) {
|
|
|
138
161
|
|
|
139
162
|
faiss::gpu::GpuIndexIVFFlatConfig config;
|
|
140
163
|
config.device = opt.device;
|
|
141
|
-
config.indicesOptions =
|
|
164
|
+
config.indicesOptions =
|
|
165
|
+
useRaft ? faiss::gpu::INDICES_64_BIT : opt.indicesOpt;
|
|
142
166
|
config.flatConfig.useFloat16 = useFloat16CoarseQuantizer;
|
|
167
|
+
config.use_raft = useRaft;
|
|
143
168
|
|
|
144
169
|
faiss::gpu::GpuIndexIVFFlat gpuIndex(
|
|
145
170
|
&res, cpuIndex.d, cpuIndex.nlist, cpuIndex.metric_type, config);
|
|
@@ -163,7 +188,7 @@ void addTest(faiss::MetricType metricType, bool useFloat16CoarseQuantizer) {
|
|
|
163
188
|
}
|
|
164
189
|
}
|
|
165
190
|
|
|
166
|
-
void copyToTest(bool useFloat16CoarseQuantizer) {
|
|
191
|
+
void copyToTest(bool useFloat16CoarseQuantizer, bool useRaft) {
|
|
167
192
|
Options opt;
|
|
168
193
|
std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
|
|
169
194
|
std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
|
|
@@ -173,8 +198,10 @@ void copyToTest(bool useFloat16CoarseQuantizer) {
|
|
|
173
198
|
|
|
174
199
|
faiss::gpu::GpuIndexIVFFlatConfig config;
|
|
175
200
|
config.device = opt.device;
|
|
176
|
-
config.indicesOptions =
|
|
201
|
+
config.indicesOptions =
|
|
202
|
+
useRaft ? faiss::gpu::INDICES_64_BIT : opt.indicesOpt;
|
|
177
203
|
config.flatConfig.useFloat16 = useFloat16CoarseQuantizer;
|
|
204
|
+
config.use_raft = useRaft;
|
|
178
205
|
|
|
179
206
|
faiss::gpu::GpuIndexIVFFlat gpuIndex(
|
|
180
207
|
&res, opt.dim, opt.numCentroids, faiss::METRIC_L2, config);
|
|
@@ -214,7 +241,7 @@ void copyToTest(bool useFloat16CoarseQuantizer) {
|
|
|
214
241
|
compFloat16 ? 0.30f : 0.015f);
|
|
215
242
|
}
|
|
216
243
|
|
|
217
|
-
void copyFromTest(bool useFloat16CoarseQuantizer) {
|
|
244
|
+
void copyFromTest(bool useFloat16CoarseQuantizer, bool useRaft) {
|
|
218
245
|
Options opt;
|
|
219
246
|
std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
|
|
220
247
|
std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
|
|
@@ -232,8 +259,10 @@ void copyFromTest(bool useFloat16CoarseQuantizer) {
|
|
|
232
259
|
|
|
233
260
|
faiss::gpu::GpuIndexIVFFlatConfig config;
|
|
234
261
|
config.device = opt.device;
|
|
235
|
-
config.indicesOptions =
|
|
262
|
+
config.indicesOptions =
|
|
263
|
+
useRaft ? faiss::gpu::INDICES_64_BIT : opt.indicesOpt;
|
|
236
264
|
config.flatConfig.useFloat16 = useFloat16CoarseQuantizer;
|
|
265
|
+
config.use_raft = useRaft;
|
|
237
266
|
|
|
238
267
|
faiss::gpu::GpuIndexIVFFlat gpuIndex(&res, 1, 1, faiss::METRIC_L2, config);
|
|
239
268
|
gpuIndex.nprobe = 1;
|
|
@@ -265,19 +294,35 @@ void copyFromTest(bool useFloat16CoarseQuantizer) {
|
|
|
265
294
|
}
|
|
266
295
|
|
|
267
296
|
TEST(TestGpuIndexIVFFlat, Float32_32_Add_L2) {
|
|
268
|
-
addTest(faiss::METRIC_L2, false);
|
|
297
|
+
addTest(faiss::METRIC_L2, false, false);
|
|
298
|
+
|
|
299
|
+
#if defined USE_NVIDIA_RAFT
|
|
300
|
+
addTest(faiss::METRIC_L2, false, true);
|
|
301
|
+
#endif
|
|
269
302
|
}
|
|
270
303
|
|
|
271
304
|
TEST(TestGpuIndexIVFFlat, Float32_32_Add_IP) {
|
|
272
|
-
addTest(faiss::METRIC_INNER_PRODUCT, false);
|
|
305
|
+
addTest(faiss::METRIC_INNER_PRODUCT, false, false);
|
|
306
|
+
|
|
307
|
+
#if defined USE_NVIDIA_RAFT
|
|
308
|
+
addTest(faiss::METRIC_INNER_PRODUCT, false, true);
|
|
309
|
+
#endif
|
|
273
310
|
}
|
|
274
311
|
|
|
275
312
|
TEST(TestGpuIndexIVFFlat, Float16_32_Add_L2) {
|
|
276
|
-
addTest(faiss::METRIC_L2, true);
|
|
313
|
+
addTest(faiss::METRIC_L2, true, false);
|
|
314
|
+
|
|
315
|
+
#if defined USE_NVIDIA_RAFT
|
|
316
|
+
addTest(faiss::METRIC_L2, true, true);
|
|
317
|
+
#endif
|
|
277
318
|
}
|
|
278
319
|
|
|
279
320
|
TEST(TestGpuIndexIVFFlat, Float16_32_Add_IP) {
|
|
280
|
-
addTest(faiss::METRIC_INNER_PRODUCT, true);
|
|
321
|
+
addTest(faiss::METRIC_INNER_PRODUCT, true, false);
|
|
322
|
+
|
|
323
|
+
#if defined USE_NVIDIA_RAFT
|
|
324
|
+
addTest(faiss::METRIC_INNER_PRODUCT, true, true);
|
|
325
|
+
#endif
|
|
281
326
|
}
|
|
282
327
|
|
|
283
328
|
//
|
|
@@ -285,11 +330,25 @@ TEST(TestGpuIndexIVFFlat, Float16_32_Add_IP) {
|
|
|
285
330
|
//
|
|
286
331
|
|
|
287
332
|
TEST(TestGpuIndexIVFFlat, Float32_Query_L2) {
|
|
288
|
-
|
|
333
|
+
Options opt;
|
|
334
|
+
queryTest(opt, faiss::METRIC_L2, false);
|
|
335
|
+
|
|
336
|
+
#if defined USE_NVIDIA_RAFT
|
|
337
|
+
opt.useRaft = true;
|
|
338
|
+
opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
|
|
339
|
+
queryTest(opt, faiss::METRIC_L2, false);
|
|
340
|
+
#endif
|
|
289
341
|
}
|
|
290
342
|
|
|
291
343
|
TEST(TestGpuIndexIVFFlat, Float32_Query_IP) {
|
|
292
|
-
|
|
344
|
+
Options opt;
|
|
345
|
+
queryTest(opt, faiss::METRIC_INNER_PRODUCT, false);
|
|
346
|
+
|
|
347
|
+
#if defined USE_NVIDIA_RAFT
|
|
348
|
+
opt.useRaft = true;
|
|
349
|
+
opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
|
|
350
|
+
queryTest(opt, faiss::METRIC_INNER_PRODUCT, false);
|
|
351
|
+
#endif
|
|
293
352
|
}
|
|
294
353
|
|
|
295
354
|
TEST(TestGpuIndexIVFFlat, LargeBatch) {
|
|
@@ -297,16 +356,36 @@ TEST(TestGpuIndexIVFFlat, LargeBatch) {
|
|
|
297
356
|
opt.dim = 3;
|
|
298
357
|
opt.numQuery = 100000;
|
|
299
358
|
queryTest(opt, faiss::METRIC_L2, false);
|
|
359
|
+
|
|
360
|
+
#if defined USE_NVIDIA_RAFT
|
|
361
|
+
opt.useRaft = true;
|
|
362
|
+
opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
|
|
363
|
+
queryTest(opt, faiss::METRIC_L2, false);
|
|
364
|
+
#endif
|
|
300
365
|
}
|
|
301
366
|
|
|
302
367
|
// float16 coarse quantizer
|
|
303
368
|
|
|
304
369
|
TEST(TestGpuIndexIVFFlat, Float16_32_Query_L2) {
|
|
305
|
-
|
|
370
|
+
Options opt;
|
|
371
|
+
queryTest(opt, faiss::METRIC_L2, true);
|
|
372
|
+
|
|
373
|
+
#if defined USE_NVIDIA_RAFT
|
|
374
|
+
opt.useRaft = true;
|
|
375
|
+
opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
|
|
376
|
+
queryTest(opt, faiss::METRIC_L2, true);
|
|
377
|
+
#endif
|
|
306
378
|
}
|
|
307
379
|
|
|
308
380
|
TEST(TestGpuIndexIVFFlat, Float16_32_Query_IP) {
|
|
309
|
-
|
|
381
|
+
Options opt;
|
|
382
|
+
queryTest(opt, faiss::METRIC_INNER_PRODUCT, true);
|
|
383
|
+
|
|
384
|
+
#if defined USE_NVIDIA_RAFT
|
|
385
|
+
opt.useRaft = true;
|
|
386
|
+
opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
|
|
387
|
+
queryTest(opt, faiss::METRIC_INNER_PRODUCT, true);
|
|
388
|
+
#endif
|
|
310
389
|
}
|
|
311
390
|
|
|
312
391
|
//
|
|
@@ -318,24 +397,48 @@ TEST(TestGpuIndexIVFFlat, Float32_Query_L2_64) {
|
|
|
318
397
|
Options opt;
|
|
319
398
|
opt.dim = 64;
|
|
320
399
|
queryTest(opt, faiss::METRIC_L2, false);
|
|
400
|
+
|
|
401
|
+
#if defined USE_NVIDIA_RAFT
|
|
402
|
+
opt.useRaft = true;
|
|
403
|
+
opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
|
|
404
|
+
queryTest(opt, faiss::METRIC_L2, false);
|
|
405
|
+
#endif
|
|
321
406
|
}
|
|
322
407
|
|
|
323
408
|
TEST(TestGpuIndexIVFFlat, Float32_Query_IP_64) {
|
|
324
409
|
Options opt;
|
|
325
410
|
opt.dim = 64;
|
|
326
411
|
queryTest(opt, faiss::METRIC_INNER_PRODUCT, false);
|
|
412
|
+
|
|
413
|
+
#if defined USE_NVIDIA_RAFT
|
|
414
|
+
opt.useRaft = true;
|
|
415
|
+
opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
|
|
416
|
+
queryTest(opt, faiss::METRIC_INNER_PRODUCT, false);
|
|
417
|
+
#endif
|
|
327
418
|
}
|
|
328
419
|
|
|
329
420
|
TEST(TestGpuIndexIVFFlat, Float32_Query_L2_128) {
|
|
330
421
|
Options opt;
|
|
331
422
|
opt.dim = 128;
|
|
332
423
|
queryTest(opt, faiss::METRIC_L2, false);
|
|
424
|
+
|
|
425
|
+
#if defined USE_NVIDIA_RAFT
|
|
426
|
+
opt.useRaft = true;
|
|
427
|
+
opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
|
|
428
|
+
queryTest(opt, faiss::METRIC_L2, false);
|
|
429
|
+
#endif
|
|
333
430
|
}
|
|
334
431
|
|
|
335
432
|
TEST(TestGpuIndexIVFFlat, Float32_Query_IP_128) {
|
|
336
433
|
Options opt;
|
|
337
434
|
opt.dim = 128;
|
|
338
435
|
queryTest(opt, faiss::METRIC_INNER_PRODUCT, false);
|
|
436
|
+
|
|
437
|
+
#if defined USE_NVIDIA_RAFT
|
|
438
|
+
opt.useRaft = true;
|
|
439
|
+
opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
|
|
440
|
+
queryTest(opt, faiss::METRIC_INNER_PRODUCT, false);
|
|
441
|
+
#endif
|
|
339
442
|
}
|
|
340
443
|
|
|
341
444
|
//
|
|
@@ -343,11 +446,19 @@ TEST(TestGpuIndexIVFFlat, Float32_Query_IP_128) {
|
|
|
343
446
|
//
|
|
344
447
|
|
|
345
448
|
TEST(TestGpuIndexIVFFlat, Float32_32_CopyTo) {
|
|
346
|
-
copyToTest(false);
|
|
449
|
+
copyToTest(false, false);
|
|
450
|
+
|
|
451
|
+
#if defined USE_NVIDIA_RAFT
|
|
452
|
+
copyToTest(false, true);
|
|
453
|
+
#endif
|
|
347
454
|
}
|
|
348
455
|
|
|
349
456
|
TEST(TestGpuIndexIVFFlat, Float32_32_CopyFrom) {
|
|
350
|
-
copyFromTest(false);
|
|
457
|
+
copyFromTest(false, false);
|
|
458
|
+
|
|
459
|
+
#if defined USE_NVIDIA_RAFT
|
|
460
|
+
copyFromTest(false, true);
|
|
461
|
+
#endif
|
|
351
462
|
}
|
|
352
463
|
|
|
353
464
|
TEST(TestGpuIndexIVFFlat, Float32_negative) {
|
|
@@ -377,23 +488,24 @@ TEST(TestGpuIndexIVFFlat, Float32_negative) {
|
|
|
377
488
|
faiss::gpu::StandardGpuResources res;
|
|
378
489
|
res.noTempMemory();
|
|
379
490
|
|
|
491
|
+
// Construct a positive test set
|
|
492
|
+
auto queryVecs = faiss::gpu::randVecs(opt.numQuery, opt.dim);
|
|
493
|
+
|
|
494
|
+
// Put all vecs on positive size
|
|
495
|
+
for (auto& f : queryVecs) {
|
|
496
|
+
f = std::abs(f);
|
|
497
|
+
}
|
|
498
|
+
|
|
380
499
|
faiss::gpu::GpuIndexIVFFlatConfig config;
|
|
381
500
|
config.device = opt.device;
|
|
382
501
|
config.indicesOptions = opt.indicesOpt;
|
|
502
|
+
config.use_raft = false;
|
|
383
503
|
|
|
384
504
|
faiss::gpu::GpuIndexIVFFlat gpuIndex(
|
|
385
505
|
&res, cpuIndex.d, cpuIndex.nlist, cpuIndex.metric_type, config);
|
|
386
506
|
gpuIndex.copyFrom(&cpuIndex);
|
|
387
507
|
gpuIndex.nprobe = opt.nprobe;
|
|
388
508
|
|
|
389
|
-
// Construct a positive test set
|
|
390
|
-
auto queryVecs = faiss::gpu::randVecs(opt.numQuery, opt.dim);
|
|
391
|
-
|
|
392
|
-
// Put all vecs on positive size
|
|
393
|
-
for (auto& f : queryVecs) {
|
|
394
|
-
f = std::abs(f);
|
|
395
|
-
}
|
|
396
|
-
|
|
397
509
|
bool compFloat16 = false;
|
|
398
510
|
faiss::gpu::compareIndices(
|
|
399
511
|
queryVecs,
|
|
@@ -409,6 +521,31 @@ TEST(TestGpuIndexIVFFlat, Float32_negative) {
|
|
|
409
521
|
// in fp16. Figure out another way to test
|
|
410
522
|
compFloat16 ? 0.99f : 0.1f,
|
|
411
523
|
compFloat16 ? 0.65f : 0.015f);
|
|
524
|
+
|
|
525
|
+
#if defined USE_NVIDIA_RAFT
|
|
526
|
+
config.use_raft = true;
|
|
527
|
+
config.indicesOptions = faiss::gpu::INDICES_64_BIT;
|
|
528
|
+
|
|
529
|
+
faiss::gpu::GpuIndexIVFFlat raftGpuIndex(
|
|
530
|
+
&res, cpuIndex.d, cpuIndex.nlist, cpuIndex.metric_type, config);
|
|
531
|
+
raftGpuIndex.copyFrom(&cpuIndex);
|
|
532
|
+
raftGpuIndex.nprobe = opt.nprobe;
|
|
533
|
+
|
|
534
|
+
faiss::gpu::compareIndices(
|
|
535
|
+
queryVecs,
|
|
536
|
+
cpuIndex,
|
|
537
|
+
raftGpuIndex,
|
|
538
|
+
opt.numQuery,
|
|
539
|
+
opt.dim,
|
|
540
|
+
opt.k,
|
|
541
|
+
opt.toString(),
|
|
542
|
+
compFloat16 ? kF16MaxRelErr : kF32MaxRelErr,
|
|
543
|
+
// FIXME: the fp16 bounds are
|
|
544
|
+
// useless when math (the accumulator) is
|
|
545
|
+
// in fp16. Figure out another way to test
|
|
546
|
+
compFloat16 ? 0.99f : 0.1f,
|
|
547
|
+
compFloat16 ? 0.65f : 0.015f);
|
|
548
|
+
#endif
|
|
412
549
|
}
|
|
413
550
|
|
|
414
551
|
//
|
|
@@ -424,10 +561,18 @@ TEST(TestGpuIndexIVFFlat, QueryNaN) {
|
|
|
424
561
|
faiss::gpu::StandardGpuResources res;
|
|
425
562
|
res.noTempMemory();
|
|
426
563
|
|
|
564
|
+
int numQuery = 10;
|
|
565
|
+
std::vector<float> nans(
|
|
566
|
+
numQuery * opt.dim, std::numeric_limits<float>::quiet_NaN());
|
|
567
|
+
|
|
568
|
+
std::vector<float> distances(numQuery * opt.k, 0);
|
|
569
|
+
std::vector<faiss::idx_t> indices(numQuery * opt.k, 0);
|
|
570
|
+
|
|
427
571
|
faiss::gpu::GpuIndexIVFFlatConfig config;
|
|
428
572
|
config.device = opt.device;
|
|
429
573
|
config.indicesOptions = opt.indicesOpt;
|
|
430
574
|
config.flatConfig.useFloat16 = faiss::gpu::randBool();
|
|
575
|
+
config.use_raft = false;
|
|
431
576
|
|
|
432
577
|
faiss::gpu::GpuIndexIVFFlat gpuIndex(
|
|
433
578
|
&res, opt.dim, opt.numCentroids, faiss::METRIC_L2, config);
|
|
@@ -436,14 +581,31 @@ TEST(TestGpuIndexIVFFlat, QueryNaN) {
|
|
|
436
581
|
gpuIndex.train(opt.numTrain, trainVecs.data());
|
|
437
582
|
gpuIndex.add(opt.numAdd, addVecs.data());
|
|
438
583
|
|
|
439
|
-
|
|
440
|
-
|
|
441
|
-
numQuery * opt.dim, std::numeric_limits<float>::quiet_NaN());
|
|
584
|
+
gpuIndex.search(
|
|
585
|
+
numQuery, nans.data(), opt.k, distances.data(), indices.data());
|
|
442
586
|
|
|
443
|
-
|
|
444
|
-
|
|
587
|
+
for (int q = 0; q < numQuery; ++q) {
|
|
588
|
+
for (int k = 0; k < opt.k; ++k) {
|
|
589
|
+
EXPECT_EQ(indices[q * opt.k + k], -1);
|
|
590
|
+
EXPECT_EQ(
|
|
591
|
+
distances[q * opt.k + k],
|
|
592
|
+
std::numeric_limits<float>::max());
|
|
593
|
+
}
|
|
594
|
+
}
|
|
445
595
|
|
|
446
|
-
|
|
596
|
+
#if defined USE_NVIDIA_RAFT
|
|
597
|
+
config.use_raft = true;
|
|
598
|
+
config.indicesOptions = faiss::gpu::INDICES_64_BIT;
|
|
599
|
+
std::fill(distances.begin(), distances.end(), 0);
|
|
600
|
+
std::fill(indices.begin(), indices.end(), 0);
|
|
601
|
+
faiss::gpu::GpuIndexIVFFlat raftGpuIndex(
|
|
602
|
+
&res, opt.dim, opt.numCentroids, faiss::METRIC_L2, config);
|
|
603
|
+
raftGpuIndex.nprobe = opt.nprobe;
|
|
604
|
+
|
|
605
|
+
raftGpuIndex.train(opt.numTrain, trainVecs.data());
|
|
606
|
+
raftGpuIndex.add(opt.numAdd, addVecs.data());
|
|
607
|
+
|
|
608
|
+
raftGpuIndex.search(
|
|
447
609
|
numQuery, nans.data(), opt.k, distances.data(), indices.data());
|
|
448
610
|
|
|
449
611
|
for (int q = 0; q < numQuery; ++q) {
|
|
@@ -454,6 +616,7 @@ TEST(TestGpuIndexIVFFlat, QueryNaN) {
|
|
|
454
616
|
std::numeric_limits<float>::max());
|
|
455
617
|
}
|
|
456
618
|
}
|
|
619
|
+
#endif
|
|
457
620
|
}
|
|
458
621
|
|
|
459
622
|
TEST(TestGpuIndexIVFFlat, AddNaN) {
|
|
@@ -462,15 +625,6 @@ TEST(TestGpuIndexIVFFlat, AddNaN) {
|
|
|
462
625
|
faiss::gpu::StandardGpuResources res;
|
|
463
626
|
res.noTempMemory();
|
|
464
627
|
|
|
465
|
-
faiss::gpu::GpuIndexIVFFlatConfig config;
|
|
466
|
-
config.device = opt.device;
|
|
467
|
-
config.indicesOptions = opt.indicesOpt;
|
|
468
|
-
config.flatConfig.useFloat16 = faiss::gpu::randBool();
|
|
469
|
-
|
|
470
|
-
faiss::gpu::GpuIndexIVFFlat gpuIndex(
|
|
471
|
-
&res, opt.dim, opt.numCentroids, faiss::METRIC_L2, config);
|
|
472
|
-
gpuIndex.nprobe = opt.nprobe;
|
|
473
|
-
|
|
474
628
|
int numNans = 10;
|
|
475
629
|
std::vector<float> nans(
|
|
476
630
|
numNans * opt.dim, std::numeric_limits<float>::quiet_NaN());
|
|
@@ -482,6 +636,15 @@ TEST(TestGpuIndexIVFFlat, AddNaN) {
|
|
|
482
636
|
}
|
|
483
637
|
|
|
484
638
|
std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
|
|
639
|
+
|
|
640
|
+
faiss::gpu::GpuIndexIVFFlatConfig config;
|
|
641
|
+
config.device = opt.device;
|
|
642
|
+
config.indicesOptions = opt.indicesOpt;
|
|
643
|
+
config.flatConfig.useFloat16 = faiss::gpu::randBool();
|
|
644
|
+
config.use_raft = false;
|
|
645
|
+
faiss::gpu::GpuIndexIVFFlat gpuIndex(
|
|
646
|
+
&res, opt.dim, opt.numCentroids, faiss::METRIC_L2, config);
|
|
647
|
+
gpuIndex.nprobe = opt.nprobe;
|
|
485
648
|
gpuIndex.train(opt.numTrain, trainVecs.data());
|
|
486
649
|
|
|
487
650
|
// should not crash
|
|
@@ -499,6 +662,27 @@ TEST(TestGpuIndexIVFFlat, AddNaN) {
|
|
|
499
662
|
opt.k,
|
|
500
663
|
distance.data(),
|
|
501
664
|
indices.data());
|
|
665
|
+
|
|
666
|
+
#if defined USE_NVIDIA_RAFT
|
|
667
|
+
config.use_raft = true;
|
|
668
|
+
config.indicesOptions = faiss::gpu::INDICES_64_BIT;
|
|
669
|
+
faiss::gpu::GpuIndexIVFFlat raftGpuIndex(
|
|
670
|
+
&res, opt.dim, opt.numCentroids, faiss::METRIC_L2, config);
|
|
671
|
+
raftGpuIndex.nprobe = opt.nprobe;
|
|
672
|
+
raftGpuIndex.train(opt.numTrain, trainVecs.data());
|
|
673
|
+
|
|
674
|
+
// should not crash
|
|
675
|
+
EXPECT_EQ(raftGpuIndex.ntotal, 0);
|
|
676
|
+
raftGpuIndex.add(numNans, nans.data());
|
|
677
|
+
|
|
678
|
+
// should not crash
|
|
679
|
+
raftGpuIndex.search(
|
|
680
|
+
opt.numQuery,
|
|
681
|
+
queryVecs.data(),
|
|
682
|
+
opt.k,
|
|
683
|
+
distance.data(),
|
|
684
|
+
indices.data());
|
|
685
|
+
#endif
|
|
502
686
|
}
|
|
503
687
|
|
|
504
688
|
TEST(TestGpuIndexIVFFlat, UnifiedMemory) {
|
|
@@ -539,6 +723,7 @@ TEST(TestGpuIndexIVFFlat, UnifiedMemory) {
|
|
|
539
723
|
faiss::gpu::GpuIndexIVFFlatConfig config;
|
|
540
724
|
config.device = device;
|
|
541
725
|
config.memorySpace = faiss::gpu::MemorySpace::Unified;
|
|
726
|
+
config.use_raft = false;
|
|
542
727
|
|
|
543
728
|
faiss::gpu::GpuIndexIVFFlat gpuIndex(
|
|
544
729
|
&res, dim, numCentroids, faiss::METRIC_L2, config);
|
|
@@ -555,6 +740,26 @@ TEST(TestGpuIndexIVFFlat, UnifiedMemory) {
|
|
|
555
740
|
kF32MaxRelErr,
|
|
556
741
|
0.1f,
|
|
557
742
|
0.015f);
|
|
743
|
+
|
|
744
|
+
#if defined USE_NVIDIA_RAFT
|
|
745
|
+
config.use_raft = true;
|
|
746
|
+
config.indicesOptions = faiss::gpu::INDICES_64_BIT;
|
|
747
|
+
faiss::gpu::GpuIndexIVFFlat raftGpuIndex(
|
|
748
|
+
&res, dim, numCentroids, faiss::METRIC_L2, config);
|
|
749
|
+
raftGpuIndex.copyFrom(&cpuIndex);
|
|
750
|
+
raftGpuIndex.nprobe = nprobe;
|
|
751
|
+
|
|
752
|
+
faiss::gpu::compareIndices(
|
|
753
|
+
cpuIndex,
|
|
754
|
+
raftGpuIndex,
|
|
755
|
+
numQuery,
|
|
756
|
+
dim,
|
|
757
|
+
k,
|
|
758
|
+
"Unified Memory",
|
|
759
|
+
kF32MaxRelErr,
|
|
760
|
+
0.1f,
|
|
761
|
+
0.015f);
|
|
762
|
+
#endif
|
|
558
763
|
}
|
|
559
764
|
|
|
560
765
|
TEST(TestGpuIndexIVFFlat, LongIVFList) {
|
|
@@ -596,6 +801,7 @@ TEST(TestGpuIndexIVFFlat, LongIVFList) {
|
|
|
596
801
|
|
|
597
802
|
faiss::gpu::GpuIndexIVFFlatConfig config;
|
|
598
803
|
config.device = device;
|
|
804
|
+
config.use_raft = false;
|
|
599
805
|
|
|
600
806
|
faiss::gpu::GpuIndexIVFFlat gpuIndex(
|
|
601
807
|
&res, dim, numCentroids, faiss::METRIC_L2, config);
|
|
@@ -613,6 +819,92 @@ TEST(TestGpuIndexIVFFlat, LongIVFList) {
|
|
|
613
819
|
kF32MaxRelErr,
|
|
614
820
|
0.1f,
|
|
615
821
|
0.015f);
|
|
822
|
+
|
|
823
|
+
#if defined USE_NVIDIA_RAFT
|
|
824
|
+
config.use_raft = true;
|
|
825
|
+
config.indicesOptions = faiss::gpu::INDICES_64_BIT;
|
|
826
|
+
faiss::gpu::GpuIndexIVFFlat raftGpuIndex(
|
|
827
|
+
&res, dim, numCentroids, faiss::METRIC_L2, config);
|
|
828
|
+
raftGpuIndex.train(numTrain, trainVecs.data());
|
|
829
|
+
raftGpuIndex.add(numAdd, addVecs.data());
|
|
830
|
+
raftGpuIndex.nprobe = 1;
|
|
831
|
+
|
|
832
|
+
faiss::gpu::compareIndices(
|
|
833
|
+
cpuIndex,
|
|
834
|
+
raftGpuIndex,
|
|
835
|
+
numQuery,
|
|
836
|
+
dim,
|
|
837
|
+
k,
|
|
838
|
+
"Unified Memory",
|
|
839
|
+
kF32MaxRelErr,
|
|
840
|
+
0.1f,
|
|
841
|
+
0.015f);
|
|
842
|
+
#endif
|
|
843
|
+
}
|
|
844
|
+
|
|
845
|
+
TEST(TestGpuIndexIVFFlat, Reconstruct_n) {
|
|
846
|
+
Options opt;
|
|
847
|
+
|
|
848
|
+
std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
|
|
849
|
+
std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
|
|
850
|
+
|
|
851
|
+
faiss::IndexFlatL2 cpuQuantizer(opt.dim);
|
|
852
|
+
faiss::IndexIVFFlat cpuIndex(
|
|
853
|
+
&cpuQuantizer, opt.dim, opt.numCentroids, faiss::METRIC_L2);
|
|
854
|
+
cpuIndex.nprobe = opt.nprobe;
|
|
855
|
+
cpuIndex.train(opt.numTrain, trainVecs.data());
|
|
856
|
+
cpuIndex.add(opt.numAdd, addVecs.data());
|
|
857
|
+
|
|
858
|
+
faiss::gpu::StandardGpuResources res;
|
|
859
|
+
res.noTempMemory();
|
|
860
|
+
|
|
861
|
+
faiss::gpu::GpuIndexIVFFlatConfig config;
|
|
862
|
+
config.device = opt.device;
|
|
863
|
+
config.indicesOptions = faiss::gpu::INDICES_64_BIT;
|
|
864
|
+
config.use_raft = false;
|
|
865
|
+
|
|
866
|
+
faiss::gpu::GpuIndexIVFFlat gpuIndex(
|
|
867
|
+
&res, opt.dim, opt.numCentroids, faiss::METRIC_L2, config);
|
|
868
|
+
gpuIndex.nprobe = opt.nprobe;
|
|
869
|
+
|
|
870
|
+
gpuIndex.train(opt.numTrain, trainVecs.data());
|
|
871
|
+
gpuIndex.add(opt.numAdd, addVecs.data());
|
|
872
|
+
|
|
873
|
+
std::vector<float> gpuVals(opt.numAdd * opt.dim);
|
|
874
|
+
|
|
875
|
+
gpuIndex.reconstruct_n(0, gpuIndex.ntotal, gpuVals.data());
|
|
876
|
+
|
|
877
|
+
std::vector<float> cpuVals(opt.numAdd * opt.dim);
|
|
878
|
+
|
|
879
|
+
cpuIndex.reconstruct_n(0, cpuIndex.ntotal, cpuVals.data());
|
|
880
|
+
|
|
881
|
+
EXPECT_EQ(gpuVals, cpuVals);
|
|
882
|
+
|
|
883
|
+
config.indicesOptions = faiss::gpu::INDICES_32_BIT;
|
|
884
|
+
|
|
885
|
+
faiss::gpu::GpuIndexIVFFlat gpuIndex1(
|
|
886
|
+
&res, opt.dim, opt.numCentroids, faiss::METRIC_L2, config);
|
|
887
|
+
gpuIndex1.nprobe = opt.nprobe;
|
|
888
|
+
|
|
889
|
+
gpuIndex1.train(opt.numTrain, trainVecs.data());
|
|
890
|
+
gpuIndex1.add(opt.numAdd, addVecs.data());
|
|
891
|
+
|
|
892
|
+
gpuIndex1.reconstruct_n(0, gpuIndex1.ntotal, gpuVals.data());
|
|
893
|
+
|
|
894
|
+
EXPECT_EQ(gpuVals, cpuVals);
|
|
895
|
+
|
|
896
|
+
config.indicesOptions = faiss::gpu::INDICES_CPU;
|
|
897
|
+
|
|
898
|
+
faiss::gpu::GpuIndexIVFFlat gpuIndex2(
|
|
899
|
+
&res, opt.dim, opt.numCentroids, faiss::METRIC_L2, config);
|
|
900
|
+
gpuIndex2.nprobe = opt.nprobe;
|
|
901
|
+
|
|
902
|
+
gpuIndex2.train(opt.numTrain, trainVecs.data());
|
|
903
|
+
gpuIndex2.add(opt.numAdd, addVecs.data());
|
|
904
|
+
|
|
905
|
+
gpuIndex2.reconstruct_n(0, gpuIndex2.ntotal, gpuVals.data());
|
|
906
|
+
|
|
907
|
+
EXPECT_EQ(gpuVals, cpuVals);
|
|
616
908
|
}
|
|
617
909
|
|
|
618
910
|
int main(int argc, char** argv) {
|