faiss 0.2.0 → 0.2.4
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 +16 -0
- data/LICENSE.txt +1 -1
- data/README.md +7 -7
- data/ext/faiss/extconf.rb +6 -3
- data/ext/faiss/numo.hpp +4 -4
- data/ext/faiss/utils.cpp +1 -1
- data/ext/faiss/utils.h +1 -1
- data/lib/faiss/version.rb +1 -1
- data/vendor/faiss/faiss/AutoTune.cpp +292 -291
- data/vendor/faiss/faiss/AutoTune.h +55 -56
- data/vendor/faiss/faiss/Clustering.cpp +365 -194
- data/vendor/faiss/faiss/Clustering.h +102 -35
- data/vendor/faiss/faiss/IVFlib.cpp +171 -195
- data/vendor/faiss/faiss/IVFlib.h +48 -51
- data/vendor/faiss/faiss/Index.cpp +85 -103
- data/vendor/faiss/faiss/Index.h +54 -48
- data/vendor/faiss/faiss/Index2Layer.cpp +126 -224
- data/vendor/faiss/faiss/Index2Layer.h +22 -36
- data/vendor/faiss/faiss/IndexAdditiveQuantizer.cpp +407 -0
- data/vendor/faiss/faiss/IndexAdditiveQuantizer.h +195 -0
- data/vendor/faiss/faiss/IndexBinary.cpp +45 -37
- data/vendor/faiss/faiss/IndexBinary.h +140 -132
- data/vendor/faiss/faiss/IndexBinaryFlat.cpp +73 -53
- data/vendor/faiss/faiss/IndexBinaryFlat.h +29 -24
- data/vendor/faiss/faiss/IndexBinaryFromFloat.cpp +46 -43
- data/vendor/faiss/faiss/IndexBinaryFromFloat.h +16 -15
- data/vendor/faiss/faiss/IndexBinaryHNSW.cpp +215 -232
- data/vendor/faiss/faiss/IndexBinaryHNSW.h +25 -24
- data/vendor/faiss/faiss/IndexBinaryHash.cpp +182 -177
- data/vendor/faiss/faiss/IndexBinaryHash.h +41 -34
- data/vendor/faiss/faiss/IndexBinaryIVF.cpp +489 -461
- data/vendor/faiss/faiss/IndexBinaryIVF.h +97 -68
- data/vendor/faiss/faiss/IndexFlat.cpp +115 -176
- data/vendor/faiss/faiss/IndexFlat.h +42 -59
- data/vendor/faiss/faiss/IndexFlatCodes.cpp +67 -0
- data/vendor/faiss/faiss/IndexFlatCodes.h +47 -0
- data/vendor/faiss/faiss/IndexHNSW.cpp +372 -348
- data/vendor/faiss/faiss/IndexHNSW.h +57 -41
- data/vendor/faiss/faiss/IndexIVF.cpp +545 -453
- data/vendor/faiss/faiss/IndexIVF.h +169 -118
- data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.cpp +316 -0
- data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.h +121 -0
- data/vendor/faiss/faiss/IndexIVFFlat.cpp +247 -252
- data/vendor/faiss/faiss/IndexIVFFlat.h +48 -51
- data/vendor/faiss/faiss/IndexIVFPQ.cpp +459 -517
- data/vendor/faiss/faiss/IndexIVFPQ.h +75 -67
- data/vendor/faiss/faiss/IndexIVFPQFastScan.cpp +406 -372
- data/vendor/faiss/faiss/IndexIVFPQFastScan.h +82 -57
- data/vendor/faiss/faiss/IndexIVFPQR.cpp +104 -102
- data/vendor/faiss/faiss/IndexIVFPQR.h +33 -28
- data/vendor/faiss/faiss/IndexIVFSpectralHash.cpp +163 -150
- data/vendor/faiss/faiss/IndexIVFSpectralHash.h +38 -25
- data/vendor/faiss/faiss/IndexLSH.cpp +66 -113
- data/vendor/faiss/faiss/IndexLSH.h +20 -38
- data/vendor/faiss/faiss/IndexLattice.cpp +42 -56
- data/vendor/faiss/faiss/IndexLattice.h +11 -16
- data/vendor/faiss/faiss/IndexNNDescent.cpp +229 -0
- data/vendor/faiss/faiss/IndexNNDescent.h +72 -0
- data/vendor/faiss/faiss/IndexNSG.cpp +301 -0
- data/vendor/faiss/faiss/IndexNSG.h +85 -0
- data/vendor/faiss/faiss/IndexPQ.cpp +387 -495
- data/vendor/faiss/faiss/IndexPQ.h +64 -82
- data/vendor/faiss/faiss/IndexPQFastScan.cpp +143 -170
- data/vendor/faiss/faiss/IndexPQFastScan.h +46 -32
- data/vendor/faiss/faiss/IndexPreTransform.cpp +120 -150
- data/vendor/faiss/faiss/IndexPreTransform.h +33 -36
- data/vendor/faiss/faiss/IndexRefine.cpp +139 -127
- data/vendor/faiss/faiss/IndexRefine.h +32 -23
- data/vendor/faiss/faiss/IndexReplicas.cpp +147 -153
- data/vendor/faiss/faiss/IndexReplicas.h +62 -56
- data/vendor/faiss/faiss/IndexScalarQuantizer.cpp +111 -172
- data/vendor/faiss/faiss/IndexScalarQuantizer.h +41 -59
- data/vendor/faiss/faiss/IndexShards.cpp +256 -240
- data/vendor/faiss/faiss/IndexShards.h +85 -73
- data/vendor/faiss/faiss/MatrixStats.cpp +112 -97
- data/vendor/faiss/faiss/MatrixStats.h +7 -10
- data/vendor/faiss/faiss/MetaIndexes.cpp +135 -157
- data/vendor/faiss/faiss/MetaIndexes.h +40 -34
- data/vendor/faiss/faiss/MetricType.h +7 -7
- data/vendor/faiss/faiss/VectorTransform.cpp +654 -475
- data/vendor/faiss/faiss/VectorTransform.h +64 -89
- data/vendor/faiss/faiss/clone_index.cpp +78 -73
- data/vendor/faiss/faiss/clone_index.h +4 -9
- data/vendor/faiss/faiss/gpu/GpuAutoTune.cpp +33 -38
- data/vendor/faiss/faiss/gpu/GpuAutoTune.h +11 -9
- data/vendor/faiss/faiss/gpu/GpuCloner.cpp +198 -171
- data/vendor/faiss/faiss/gpu/GpuCloner.h +53 -35
- data/vendor/faiss/faiss/gpu/GpuClonerOptions.cpp +12 -14
- data/vendor/faiss/faiss/gpu/GpuClonerOptions.h +27 -25
- data/vendor/faiss/faiss/gpu/GpuDistance.h +116 -112
- data/vendor/faiss/faiss/gpu/GpuFaissAssert.h +1 -2
- data/vendor/faiss/faiss/gpu/GpuIcmEncoder.h +60 -0
- data/vendor/faiss/faiss/gpu/GpuIndex.h +134 -137
- data/vendor/faiss/faiss/gpu/GpuIndexBinaryFlat.h +76 -73
- data/vendor/faiss/faiss/gpu/GpuIndexFlat.h +173 -162
- data/vendor/faiss/faiss/gpu/GpuIndexIVF.h +67 -64
- data/vendor/faiss/faiss/gpu/GpuIndexIVFFlat.h +89 -86
- data/vendor/faiss/faiss/gpu/GpuIndexIVFPQ.h +150 -141
- data/vendor/faiss/faiss/gpu/GpuIndexIVFScalarQuantizer.h +101 -103
- data/vendor/faiss/faiss/gpu/GpuIndicesOptions.h +17 -16
- data/vendor/faiss/faiss/gpu/GpuResources.cpp +116 -128
- data/vendor/faiss/faiss/gpu/GpuResources.h +182 -186
- data/vendor/faiss/faiss/gpu/StandardGpuResources.cpp +433 -422
- data/vendor/faiss/faiss/gpu/StandardGpuResources.h +131 -130
- data/vendor/faiss/faiss/gpu/impl/InterleavedCodes.cpp +468 -456
- data/vendor/faiss/faiss/gpu/impl/InterleavedCodes.h +25 -19
- data/vendor/faiss/faiss/gpu/impl/RemapIndices.cpp +22 -20
- data/vendor/faiss/faiss/gpu/impl/RemapIndices.h +9 -8
- data/vendor/faiss/faiss/gpu/perf/IndexWrapper-inl.h +39 -44
- data/vendor/faiss/faiss/gpu/perf/IndexWrapper.h +16 -14
- data/vendor/faiss/faiss/gpu/perf/PerfClustering.cpp +77 -71
- data/vendor/faiss/faiss/gpu/perf/PerfIVFPQAdd.cpp +109 -88
- data/vendor/faiss/faiss/gpu/perf/WriteIndex.cpp +75 -64
- data/vendor/faiss/faiss/gpu/test/TestCodePacking.cpp +230 -215
- data/vendor/faiss/faiss/gpu/test/TestGpuIndexBinaryFlat.cpp +80 -86
- data/vendor/faiss/faiss/gpu/test/TestGpuIndexFlat.cpp +284 -277
- data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFFlat.cpp +416 -416
- data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFPQ.cpp +611 -517
- data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFScalarQuantizer.cpp +166 -164
- data/vendor/faiss/faiss/gpu/test/TestGpuMemoryException.cpp +61 -53
- data/vendor/faiss/faiss/gpu/test/TestUtils.cpp +274 -238
- data/vendor/faiss/faiss/gpu/test/TestUtils.h +73 -57
- data/vendor/faiss/faiss/gpu/test/demo_ivfpq_indexing_gpu.cpp +47 -50
- data/vendor/faiss/faiss/gpu/utils/DeviceUtils.h +79 -72
- data/vendor/faiss/faiss/gpu/utils/StackDeviceMemory.cpp +140 -146
- data/vendor/faiss/faiss/gpu/utils/StackDeviceMemory.h +69 -71
- data/vendor/faiss/faiss/gpu/utils/StaticUtils.h +21 -16
- data/vendor/faiss/faiss/gpu/utils/Timer.cpp +25 -29
- data/vendor/faiss/faiss/gpu/utils/Timer.h +30 -29
- data/vendor/faiss/faiss/impl/AdditiveQuantizer.cpp +503 -0
- data/vendor/faiss/faiss/impl/AdditiveQuantizer.h +175 -0
- data/vendor/faiss/faiss/impl/AuxIndexStructures.cpp +90 -120
- data/vendor/faiss/faiss/impl/AuxIndexStructures.h +81 -65
- data/vendor/faiss/faiss/impl/FaissAssert.h +73 -58
- data/vendor/faiss/faiss/impl/FaissException.cpp +56 -48
- data/vendor/faiss/faiss/impl/FaissException.h +41 -29
- data/vendor/faiss/faiss/impl/HNSW.cpp +606 -617
- data/vendor/faiss/faiss/impl/HNSW.h +179 -200
- data/vendor/faiss/faiss/impl/LocalSearchQuantizer.cpp +855 -0
- data/vendor/faiss/faiss/impl/LocalSearchQuantizer.h +244 -0
- data/vendor/faiss/faiss/impl/NNDescent.cpp +487 -0
- data/vendor/faiss/faiss/impl/NNDescent.h +154 -0
- data/vendor/faiss/faiss/impl/NSG.cpp +679 -0
- data/vendor/faiss/faiss/impl/NSG.h +199 -0
- data/vendor/faiss/faiss/impl/PolysemousTraining.cpp +484 -454
- data/vendor/faiss/faiss/impl/PolysemousTraining.h +52 -55
- data/vendor/faiss/faiss/impl/ProductQuantizer-inl.h +26 -47
- data/vendor/faiss/faiss/impl/ProductQuantizer.cpp +469 -459
- data/vendor/faiss/faiss/impl/ProductQuantizer.h +76 -87
- data/vendor/faiss/faiss/impl/ResidualQuantizer.cpp +758 -0
- data/vendor/faiss/faiss/impl/ResidualQuantizer.h +188 -0
- data/vendor/faiss/faiss/impl/ResultHandler.h +96 -132
- data/vendor/faiss/faiss/impl/ScalarQuantizer.cpp +647 -707
- data/vendor/faiss/faiss/impl/ScalarQuantizer.h +48 -46
- data/vendor/faiss/faiss/impl/ThreadedIndex-inl.h +129 -131
- data/vendor/faiss/faiss/impl/ThreadedIndex.h +61 -55
- data/vendor/faiss/faiss/impl/index_read.cpp +631 -480
- data/vendor/faiss/faiss/impl/index_write.cpp +547 -407
- data/vendor/faiss/faiss/impl/io.cpp +76 -95
- data/vendor/faiss/faiss/impl/io.h +31 -41
- data/vendor/faiss/faiss/impl/io_macros.h +60 -29
- data/vendor/faiss/faiss/impl/kmeans1d.cpp +301 -0
- data/vendor/faiss/faiss/impl/kmeans1d.h +48 -0
- data/vendor/faiss/faiss/impl/lattice_Zn.cpp +137 -186
- data/vendor/faiss/faiss/impl/lattice_Zn.h +40 -51
- data/vendor/faiss/faiss/impl/platform_macros.h +29 -8
- data/vendor/faiss/faiss/impl/pq4_fast_scan.cpp +77 -124
- data/vendor/faiss/faiss/impl/pq4_fast_scan.h +39 -48
- data/vendor/faiss/faiss/impl/pq4_fast_scan_search_1.cpp +41 -52
- data/vendor/faiss/faiss/impl/pq4_fast_scan_search_qbs.cpp +80 -117
- data/vendor/faiss/faiss/impl/simd_result_handlers.h +109 -137
- data/vendor/faiss/faiss/index_factory.cpp +619 -397
- data/vendor/faiss/faiss/index_factory.h +8 -6
- data/vendor/faiss/faiss/index_io.h +23 -26
- data/vendor/faiss/faiss/invlists/BlockInvertedLists.cpp +67 -75
- data/vendor/faiss/faiss/invlists/BlockInvertedLists.h +22 -24
- data/vendor/faiss/faiss/invlists/DirectMap.cpp +96 -112
- data/vendor/faiss/faiss/invlists/DirectMap.h +29 -33
- data/vendor/faiss/faiss/invlists/InvertedLists.cpp +307 -364
- data/vendor/faiss/faiss/invlists/InvertedLists.h +151 -151
- data/vendor/faiss/faiss/invlists/InvertedListsIOHook.cpp +29 -34
- data/vendor/faiss/faiss/invlists/InvertedListsIOHook.h +17 -18
- data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.cpp +257 -293
- data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.h +50 -45
- data/vendor/faiss/faiss/python/python_callbacks.cpp +23 -26
- data/vendor/faiss/faiss/python/python_callbacks.h +9 -16
- data/vendor/faiss/faiss/utils/AlignedTable.h +79 -44
- data/vendor/faiss/faiss/utils/Heap.cpp +40 -48
- data/vendor/faiss/faiss/utils/Heap.h +186 -209
- data/vendor/faiss/faiss/utils/WorkerThread.cpp +67 -76
- data/vendor/faiss/faiss/utils/WorkerThread.h +32 -33
- data/vendor/faiss/faiss/utils/distances.cpp +305 -312
- data/vendor/faiss/faiss/utils/distances.h +170 -122
- data/vendor/faiss/faiss/utils/distances_simd.cpp +498 -508
- data/vendor/faiss/faiss/utils/extra_distances-inl.h +117 -0
- data/vendor/faiss/faiss/utils/extra_distances.cpp +113 -232
- data/vendor/faiss/faiss/utils/extra_distances.h +30 -29
- data/vendor/faiss/faiss/utils/hamming-inl.h +260 -209
- data/vendor/faiss/faiss/utils/hamming.cpp +375 -469
- data/vendor/faiss/faiss/utils/hamming.h +62 -85
- data/vendor/faiss/faiss/utils/ordered_key_value.h +16 -18
- data/vendor/faiss/faiss/utils/partitioning.cpp +393 -318
- data/vendor/faiss/faiss/utils/partitioning.h +26 -21
- data/vendor/faiss/faiss/utils/quantize_lut.cpp +78 -66
- data/vendor/faiss/faiss/utils/quantize_lut.h +22 -20
- data/vendor/faiss/faiss/utils/random.cpp +39 -63
- data/vendor/faiss/faiss/utils/random.h +13 -16
- data/vendor/faiss/faiss/utils/simdlib.h +4 -2
- data/vendor/faiss/faiss/utils/simdlib_avx2.h +88 -85
- data/vendor/faiss/faiss/utils/simdlib_emulated.h +226 -165
- data/vendor/faiss/faiss/utils/simdlib_neon.h +832 -0
- data/vendor/faiss/faiss/utils/utils.cpp +304 -287
- data/vendor/faiss/faiss/utils/utils.h +54 -49
- metadata +29 -4
|
@@ -5,20 +5,20 @@
|
|
|
5
5
|
* LICENSE file in the root directory of this source tree.
|
|
6
6
|
*/
|
|
7
7
|
|
|
8
|
-
|
|
9
8
|
#pragma once
|
|
10
9
|
|
|
11
|
-
#include <faiss/impl/FaissAssert.h>
|
|
12
10
|
#include <faiss/Index.h>
|
|
11
|
+
#include <faiss/impl/FaissAssert.h>
|
|
13
12
|
#include <faiss/invlists/InvertedLists.h>
|
|
14
|
-
#include <initializer_list>
|
|
15
13
|
#include <gtest/gtest.h>
|
|
16
14
|
#include <cstring>
|
|
15
|
+
#include <initializer_list>
|
|
17
16
|
#include <memory>
|
|
18
17
|
#include <string>
|
|
19
18
|
#include <vector>
|
|
20
19
|
|
|
21
|
-
namespace faiss {
|
|
20
|
+
namespace faiss {
|
|
21
|
+
namespace gpu {
|
|
22
22
|
|
|
23
23
|
/// Generates and displays a new seed for the test
|
|
24
24
|
void newTestSeed();
|
|
@@ -40,18 +40,18 @@ bool randBool();
|
|
|
40
40
|
/// initializer_list
|
|
41
41
|
template <typename T>
|
|
42
42
|
T randSelect(std::initializer_list<T> vals) {
|
|
43
|
-
|
|
44
|
-
|
|
45
|
-
|
|
46
|
-
|
|
47
|
-
|
|
48
|
-
|
|
49
|
-
|
|
43
|
+
FAISS_ASSERT(vals.size() > 0);
|
|
44
|
+
int sel = randVal(0, vals.size());
|
|
45
|
+
|
|
46
|
+
int i = 0;
|
|
47
|
+
for (auto v : vals) {
|
|
48
|
+
if (i++ == sel) {
|
|
49
|
+
return v;
|
|
50
|
+
}
|
|
50
51
|
}
|
|
51
|
-
}
|
|
52
52
|
|
|
53
|
-
|
|
54
|
-
|
|
53
|
+
// should not get here
|
|
54
|
+
return *vals.begin();
|
|
55
55
|
}
|
|
56
56
|
|
|
57
57
|
/// Generates a collection of random vectors in the range [0, 1]
|
|
@@ -62,65 +62,81 @@ std::vector<unsigned char> randBinaryVecs(size_t num, size_t dim);
|
|
|
62
62
|
|
|
63
63
|
/// Compare two indices via query for similarity, with a user-specified set of
|
|
64
64
|
/// query vectors
|
|
65
|
-
void compareIndices(
|
|
66
|
-
|
|
67
|
-
|
|
68
|
-
|
|
69
|
-
|
|
70
|
-
|
|
71
|
-
|
|
72
|
-
|
|
65
|
+
void compareIndices(
|
|
66
|
+
const std::vector<float>& queryVecs,
|
|
67
|
+
faiss::Index& refIndex,
|
|
68
|
+
faiss::Index& testIndex,
|
|
69
|
+
int numQuery,
|
|
70
|
+
int dim,
|
|
71
|
+
int k,
|
|
72
|
+
const std::string& configMsg,
|
|
73
|
+
float maxRelativeError = 6e-5f,
|
|
74
|
+
float pctMaxDiff1 = 0.1f,
|
|
75
|
+
float pctMaxDiffN = 0.005f);
|
|
73
76
|
|
|
74
77
|
/// Compare two indices via query for similarity, generating random query
|
|
75
78
|
/// vectors
|
|
76
|
-
void compareIndices(
|
|
77
|
-
|
|
78
|
-
|
|
79
|
-
|
|
80
|
-
|
|
81
|
-
|
|
82
|
-
|
|
79
|
+
void compareIndices(
|
|
80
|
+
faiss::Index& refIndex,
|
|
81
|
+
faiss::Index& testIndex,
|
|
82
|
+
int numQuery,
|
|
83
|
+
int dim,
|
|
84
|
+
int k,
|
|
85
|
+
const std::string& configMsg,
|
|
86
|
+
float maxRelativeError = 6e-5f,
|
|
87
|
+
float pctMaxDiff1 = 0.1f,
|
|
88
|
+
float pctMaxDiffN = 0.005f);
|
|
83
89
|
|
|
84
90
|
/// Display specific differences in the two (distance, index) lists
|
|
85
|
-
void compareLists(
|
|
86
|
-
|
|
87
|
-
|
|
88
|
-
|
|
89
|
-
|
|
90
|
-
|
|
91
|
-
|
|
92
|
-
|
|
93
|
-
|
|
94
|
-
|
|
91
|
+
void compareLists(
|
|
92
|
+
const float* refDist,
|
|
93
|
+
const faiss::Index::idx_t* refInd,
|
|
94
|
+
const float* testDist,
|
|
95
|
+
const faiss::Index::idx_t* testInd,
|
|
96
|
+
int dim1,
|
|
97
|
+
int dim2,
|
|
98
|
+
const std::string& configMsg,
|
|
99
|
+
bool printBasicStats,
|
|
100
|
+
bool printDiffs,
|
|
101
|
+
bool assertOnErr,
|
|
102
|
+
float maxRelativeError = 6e-5f,
|
|
103
|
+
float pctMaxDiff1 = 0.1f,
|
|
104
|
+
float pctMaxDiffN = 0.005f);
|
|
95
105
|
|
|
96
106
|
/// Compare IVF lists between a CPU and GPU index
|
|
97
107
|
template <typename A, typename B>
|
|
98
108
|
void testIVFEquality(A& cpuIndex, B& gpuIndex) {
|
|
99
|
-
|
|
100
|
-
|
|
109
|
+
// Ensure equality of the inverted lists
|
|
110
|
+
EXPECT_EQ(cpuIndex.nlist, gpuIndex.nlist);
|
|
101
111
|
|
|
102
|
-
|
|
103
|
-
|
|
112
|
+
for (int i = 0; i < cpuIndex.nlist; ++i) {
|
|
113
|
+
auto cpuLists = cpuIndex.invlists;
|
|
104
114
|
|
|
105
|
-
|
|
106
|
-
|
|
107
|
-
|
|
115
|
+
// Code equality
|
|
116
|
+
EXPECT_EQ(cpuLists->list_size(i), gpuIndex.getListLength(i));
|
|
117
|
+
std::vector<uint8_t> cpuCodes(
|
|
118
|
+
cpuLists->list_size(i) * cpuLists->code_size);
|
|
108
119
|
|
|
109
|
-
|
|
110
|
-
|
|
120
|
+
auto sc = faiss::InvertedLists::ScopedCodes(cpuLists, i);
|
|
121
|
+
std::memcpy(
|
|
122
|
+
cpuCodes.data(),
|
|
123
|
+
sc.get(),
|
|
111
124
|
cpuLists->list_size(i) * cpuLists->code_size);
|
|
112
125
|
|
|
113
|
-
|
|
114
|
-
|
|
126
|
+
auto gpuCodes = gpuIndex.getListVectorData(i, false);
|
|
127
|
+
EXPECT_EQ(cpuCodes, gpuCodes);
|
|
115
128
|
|
|
116
|
-
|
|
117
|
-
|
|
129
|
+
// Index equality
|
|
130
|
+
std::vector<Index::idx_t> cpuIndices(cpuLists->list_size(i));
|
|
118
131
|
|
|
119
|
-
|
|
120
|
-
|
|
132
|
+
auto si = faiss::InvertedLists::ScopedIds(cpuLists, i);
|
|
133
|
+
std::memcpy(
|
|
134
|
+
cpuIndices.data(),
|
|
135
|
+
si.get(),
|
|
121
136
|
cpuLists->list_size(i) * sizeof(faiss::Index::idx_t));
|
|
122
|
-
|
|
123
|
-
|
|
137
|
+
EXPECT_EQ(cpuIndices, gpuIndex.getListIndices(i));
|
|
138
|
+
}
|
|
124
139
|
}
|
|
125
140
|
|
|
126
|
-
}
|
|
141
|
+
} // namespace gpu
|
|
142
|
+
} // namespace faiss
|
|
@@ -12,24 +12,19 @@
|
|
|
12
12
|
|
|
13
13
|
#include <sys/time.h>
|
|
14
14
|
|
|
15
|
-
|
|
16
|
-
#include <faiss/gpu/StandardGpuResources.h>
|
|
17
15
|
#include <faiss/gpu/GpuIndexIVFPQ.h>
|
|
16
|
+
#include <faiss/gpu/StandardGpuResources.h>
|
|
18
17
|
|
|
19
18
|
#include <faiss/gpu/GpuAutoTune.h>
|
|
20
19
|
#include <faiss/index_io.h>
|
|
21
20
|
|
|
22
|
-
double elapsed
|
|
23
|
-
{
|
|
21
|
+
double elapsed() {
|
|
24
22
|
struct timeval tv;
|
|
25
|
-
gettimeofday
|
|
26
|
-
return
|
|
23
|
+
gettimeofday(&tv, NULL);
|
|
24
|
+
return tv.tv_sec + tv.tv_usec * 1e-6;
|
|
27
25
|
}
|
|
28
26
|
|
|
29
|
-
|
|
30
|
-
int main ()
|
|
31
|
-
{
|
|
32
|
-
|
|
27
|
+
int main() {
|
|
33
28
|
double t0 = elapsed();
|
|
34
29
|
|
|
35
30
|
// dimension of the vectors to index
|
|
@@ -48,47 +43,48 @@ int main ()
|
|
|
48
43
|
elapsed() - t0, d, nb, nt, dev_no);
|
|
49
44
|
*/
|
|
50
45
|
// a reasonable number of centroids to index nb vectors
|
|
51
|
-
int ncentroids = int
|
|
46
|
+
int ncentroids = int(4 * sqrt(nb));
|
|
52
47
|
|
|
53
48
|
faiss::gpu::StandardGpuResources resources;
|
|
54
49
|
|
|
55
|
-
|
|
56
50
|
// the coarse quantizer should not be dealloced before the index
|
|
57
51
|
// 4 = nb of bytes per code (d must be a multiple of this)
|
|
58
52
|
// 8 = nb of bits per sub-code (almost always 8)
|
|
59
53
|
faiss::gpu::GpuIndexIVFPQConfig config;
|
|
60
54
|
config.device = dev_no;
|
|
61
55
|
|
|
62
|
-
faiss::gpu::GpuIndexIVFPQ index
|
|
63
|
-
|
|
56
|
+
faiss::gpu::GpuIndexIVFPQ index(
|
|
57
|
+
&resources, d, ncentroids, 4, 8, faiss::METRIC_L2, config);
|
|
64
58
|
|
|
65
59
|
std::mt19937 rng;
|
|
66
60
|
|
|
67
61
|
{ // training
|
|
68
|
-
printf
|
|
69
|
-
|
|
62
|
+
printf("[%.3f s] Generating %ld vectors in %dD for training\n",
|
|
63
|
+
elapsed() - t0,
|
|
64
|
+
nt,
|
|
65
|
+
d);
|
|
70
66
|
|
|
71
|
-
std::vector
|
|
67
|
+
std::vector<float> trainvecs(nt * d);
|
|
72
68
|
std::uniform_real_distribution<> distrib;
|
|
73
69
|
for (size_t i = 0; i < nt * d; i++) {
|
|
74
70
|
trainvecs[i] = distrib(rng);
|
|
75
71
|
}
|
|
76
72
|
|
|
77
|
-
printf
|
|
78
|
-
elapsed() - t0);
|
|
73
|
+
printf("[%.3f s] Training the index\n", elapsed() - t0);
|
|
79
74
|
index.verbose = true;
|
|
80
75
|
|
|
81
|
-
index.train
|
|
76
|
+
index.train(nt, trainvecs.data());
|
|
82
77
|
}
|
|
83
78
|
|
|
84
79
|
{ // I/O demo
|
|
85
|
-
const char
|
|
86
|
-
printf
|
|
87
|
-
|
|
80
|
+
const char* outfilename = "/tmp/index_trained.faissindex";
|
|
81
|
+
printf("[%.3f s] storing the pre-trained index to %s\n",
|
|
82
|
+
elapsed() - t0,
|
|
83
|
+
outfilename);
|
|
88
84
|
|
|
89
|
-
faiss::Index
|
|
85
|
+
faiss::Index* cpu_index = faiss::gpu::index_gpu_to_cpu(&index);
|
|
90
86
|
|
|
91
|
-
write_index
|
|
87
|
+
write_index(cpu_index, outfilename);
|
|
92
88
|
|
|
93
89
|
delete cpu_index;
|
|
94
90
|
}
|
|
@@ -97,64 +93,65 @@ int main ()
|
|
|
97
93
|
std::vector<float> queries;
|
|
98
94
|
|
|
99
95
|
{ // populating the database
|
|
100
|
-
printf
|
|
101
|
-
|
|
96
|
+
printf("[%.3f s] Building a dataset of %ld vectors to index\n",
|
|
97
|
+
elapsed() - t0,
|
|
98
|
+
nb);
|
|
102
99
|
|
|
103
|
-
std::vector
|
|
100
|
+
std::vector<float> database(nb * d);
|
|
104
101
|
std::uniform_real_distribution<> distrib;
|
|
105
102
|
for (size_t i = 0; i < nb * d; i++) {
|
|
106
103
|
database[i] = distrib(rng);
|
|
107
104
|
}
|
|
108
105
|
|
|
109
|
-
printf
|
|
110
|
-
elapsed() - t0);
|
|
106
|
+
printf("[%.3f s] Adding the vectors to the index\n", elapsed() - t0);
|
|
111
107
|
|
|
112
|
-
index.add
|
|
108
|
+
index.add(nb, database.data());
|
|
113
109
|
|
|
114
|
-
printf
|
|
110
|
+
printf("[%.3f s] done\n", elapsed() - t0);
|
|
115
111
|
|
|
116
112
|
// remember a few elements from the database as queries
|
|
117
113
|
int i0 = 1234;
|
|
118
114
|
int i1 = 1243;
|
|
119
115
|
|
|
120
116
|
nq = i1 - i0;
|
|
121
|
-
queries.resize
|
|
117
|
+
queries.resize(nq * d);
|
|
122
118
|
for (int i = i0; i < i1; i++) {
|
|
123
119
|
for (int j = 0; j < d; j++) {
|
|
124
|
-
queries
|
|
120
|
+
queries[(i - i0) * d + j] = database[i * d + j];
|
|
125
121
|
}
|
|
126
122
|
}
|
|
127
|
-
|
|
128
123
|
}
|
|
129
124
|
|
|
130
125
|
{ // searching the database
|
|
131
126
|
int k = 5;
|
|
132
|
-
printf
|
|
133
|
-
|
|
134
|
-
|
|
127
|
+
printf("[%.3f s] Searching the %d nearest neighbors "
|
|
128
|
+
"of %ld vectors in the index\n",
|
|
129
|
+
elapsed() - t0,
|
|
130
|
+
k,
|
|
131
|
+
nq);
|
|
135
132
|
|
|
136
|
-
std::vector<faiss::Index::idx_t> nns
|
|
137
|
-
std::vector<float>
|
|
133
|
+
std::vector<faiss::Index::idx_t> nns(k * nq);
|
|
134
|
+
std::vector<float> dis(k * nq);
|
|
138
135
|
|
|
139
|
-
index.search
|
|
136
|
+
index.search(nq, queries.data(), k, dis.data(), nns.data());
|
|
140
137
|
|
|
141
|
-
printf
|
|
142
|
-
|
|
138
|
+
printf("[%.3f s] Query results (vector ids, then distances):\n",
|
|
139
|
+
elapsed() - t0);
|
|
143
140
|
|
|
144
141
|
for (int i = 0; i < nq; i++) {
|
|
145
|
-
printf
|
|
142
|
+
printf("query %2d: ", i);
|
|
146
143
|
for (int j = 0; j < k; j++) {
|
|
147
|
-
printf
|
|
144
|
+
printf("%7ld ", nns[j + i * k]);
|
|
148
145
|
}
|
|
149
|
-
printf
|
|
146
|
+
printf("\n dis: ");
|
|
150
147
|
for (int j = 0; j < k; j++) {
|
|
151
|
-
printf
|
|
148
|
+
printf("%7g ", dis[j + i * k]);
|
|
152
149
|
}
|
|
153
|
-
printf
|
|
150
|
+
printf("\n");
|
|
154
151
|
}
|
|
155
152
|
|
|
156
|
-
printf
|
|
157
|
-
|
|
153
|
+
printf("note that the nearest neighbor is not at "
|
|
154
|
+
"distance 0 due to quantization errors\n");
|
|
158
155
|
}
|
|
159
156
|
|
|
160
157
|
return 0;
|
|
@@ -5,15 +5,15 @@
|
|
|
5
5
|
* LICENSE file in the root directory of this source tree.
|
|
6
6
|
*/
|
|
7
7
|
|
|
8
|
-
|
|
9
8
|
#pragma once
|
|
10
9
|
|
|
11
|
-
#include <faiss/impl/FaissAssert.h>
|
|
12
|
-
#include <cuda_runtime.h>
|
|
13
10
|
#include <cublas_v2.h>
|
|
11
|
+
#include <cuda_runtime.h>
|
|
12
|
+
#include <faiss/impl/FaissAssert.h>
|
|
14
13
|
#include <vector>
|
|
15
14
|
|
|
16
|
-
namespace faiss {
|
|
15
|
+
namespace faiss {
|
|
16
|
+
namespace gpu {
|
|
17
17
|
|
|
18
18
|
/// Returns the current thread-local GPU device
|
|
19
19
|
int getCurrentDevice();
|
|
@@ -78,114 +78,121 @@ int getMaxKSelection();
|
|
|
78
78
|
/// RAII object to set the current device, and restore the previous
|
|
79
79
|
/// device upon destruction
|
|
80
80
|
class DeviceScope {
|
|
81
|
-
|
|
82
|
-
|
|
83
|
-
|
|
81
|
+
public:
|
|
82
|
+
explicit DeviceScope(int device);
|
|
83
|
+
~DeviceScope();
|
|
84
84
|
|
|
85
|
-
|
|
86
|
-
|
|
85
|
+
private:
|
|
86
|
+
int prevDevice_;
|
|
87
87
|
};
|
|
88
88
|
|
|
89
89
|
/// RAII object to manage a cublasHandle_t
|
|
90
90
|
class CublasHandleScope {
|
|
91
|
-
|
|
92
|
-
|
|
93
|
-
|
|
91
|
+
public:
|
|
92
|
+
CublasHandleScope();
|
|
93
|
+
~CublasHandleScope();
|
|
94
94
|
|
|
95
|
-
|
|
95
|
+
cublasHandle_t get() {
|
|
96
|
+
return blasHandle_;
|
|
97
|
+
}
|
|
96
98
|
|
|
97
|
-
|
|
98
|
-
|
|
99
|
+
private:
|
|
100
|
+
cublasHandle_t blasHandle_;
|
|
99
101
|
};
|
|
100
102
|
|
|
101
103
|
// RAII object to manage a cudaEvent_t
|
|
102
104
|
class CudaEvent {
|
|
103
|
-
|
|
104
|
-
|
|
105
|
-
|
|
106
|
-
|
|
107
|
-
|
|
108
|
-
|
|
109
|
-
|
|
110
|
-
|
|
105
|
+
public:
|
|
106
|
+
/// Creates an event and records it in this stream
|
|
107
|
+
explicit CudaEvent(cudaStream_t stream, bool timer = false);
|
|
108
|
+
CudaEvent(const CudaEvent& event) = delete;
|
|
109
|
+
CudaEvent(CudaEvent&& event) noexcept;
|
|
110
|
+
~CudaEvent();
|
|
111
|
+
|
|
112
|
+
inline cudaEvent_t get() {
|
|
113
|
+
return event_;
|
|
114
|
+
}
|
|
111
115
|
|
|
112
|
-
|
|
113
|
-
|
|
116
|
+
/// Wait on this event in this stream
|
|
117
|
+
void streamWaitOnEvent(cudaStream_t stream);
|
|
114
118
|
|
|
115
|
-
|
|
116
|
-
|
|
119
|
+
/// Have the CPU wait for the completion of this event
|
|
120
|
+
void cpuWaitOnEvent();
|
|
117
121
|
|
|
118
|
-
|
|
119
|
-
|
|
122
|
+
CudaEvent& operator=(CudaEvent&& event) noexcept;
|
|
123
|
+
CudaEvent& operator=(CudaEvent& event) = delete;
|
|
120
124
|
|
|
121
|
-
|
|
122
|
-
|
|
125
|
+
private:
|
|
126
|
+
cudaEvent_t event_;
|
|
123
127
|
};
|
|
124
128
|
|
|
125
129
|
/// Wrapper to test return status of CUDA functions
|
|
126
|
-
#define CUDA_VERIFY(X)
|
|
127
|
-
|
|
128
|
-
|
|
129
|
-
|
|
130
|
-
|
|
131
|
-
|
|
130
|
+
#define CUDA_VERIFY(X) \
|
|
131
|
+
do { \
|
|
132
|
+
auto err__ = (X); \
|
|
133
|
+
FAISS_ASSERT_FMT( \
|
|
134
|
+
err__ == cudaSuccess, \
|
|
135
|
+
"CUDA error %d %s", \
|
|
136
|
+
(int)err__, \
|
|
137
|
+
cudaGetErrorString(err__)); \
|
|
138
|
+
} while (0)
|
|
132
139
|
|
|
133
140
|
/// Wrapper to synchronously probe for CUDA errors
|
|
134
141
|
// #define FAISS_GPU_SYNC_ERROR 1
|
|
135
142
|
|
|
136
143
|
#ifdef FAISS_GPU_SYNC_ERROR
|
|
137
|
-
#define CUDA_TEST_ERROR()
|
|
138
|
-
|
|
139
|
-
|
|
140
|
-
|
|
144
|
+
#define CUDA_TEST_ERROR() \
|
|
145
|
+
do { \
|
|
146
|
+
CUDA_VERIFY(cudaDeviceSynchronize()); \
|
|
147
|
+
} while (0)
|
|
141
148
|
#else
|
|
142
|
-
#define CUDA_TEST_ERROR()
|
|
143
|
-
|
|
144
|
-
|
|
145
|
-
|
|
149
|
+
#define CUDA_TEST_ERROR() \
|
|
150
|
+
do { \
|
|
151
|
+
CUDA_VERIFY(cudaGetLastError()); \
|
|
152
|
+
} while (0)
|
|
146
153
|
#endif
|
|
147
154
|
|
|
148
155
|
/// Call for a collection of streams to wait on
|
|
149
156
|
template <typename L1, typename L2>
|
|
150
157
|
void streamWaitBase(const L1& listWaiting, const L2& listWaitOn) {
|
|
151
|
-
|
|
152
|
-
|
|
153
|
-
|
|
154
|
-
|
|
155
|
-
|
|
156
|
-
|
|
157
|
-
|
|
158
|
-
|
|
159
|
-
|
|
160
|
-
|
|
161
|
-
|
|
162
|
-
|
|
163
|
-
|
|
158
|
+
// For all the streams we are waiting on, create an event
|
|
159
|
+
std::vector<cudaEvent_t> events;
|
|
160
|
+
for (auto& stream : listWaitOn) {
|
|
161
|
+
cudaEvent_t event;
|
|
162
|
+
CUDA_VERIFY(cudaEventCreateWithFlags(&event, cudaEventDisableTiming));
|
|
163
|
+
CUDA_VERIFY(cudaEventRecord(event, stream));
|
|
164
|
+
events.push_back(event);
|
|
165
|
+
}
|
|
166
|
+
|
|
167
|
+
// For all the streams that are waiting, issue a wait
|
|
168
|
+
for (auto& stream : listWaiting) {
|
|
169
|
+
for (auto& event : events) {
|
|
170
|
+
CUDA_VERIFY(cudaStreamWaitEvent(stream, event, 0));
|
|
171
|
+
}
|
|
164
172
|
}
|
|
165
|
-
}
|
|
166
173
|
|
|
167
|
-
|
|
168
|
-
|
|
169
|
-
|
|
174
|
+
for (auto& event : events) {
|
|
175
|
+
CUDA_VERIFY(cudaEventDestroy(event));
|
|
176
|
+
}
|
|
170
177
|
}
|
|
171
178
|
|
|
172
179
|
/// These versions allow usage of initializer_list as arguments, since
|
|
173
180
|
/// otherwise {...} doesn't have a type
|
|
174
181
|
template <typename L1>
|
|
175
|
-
void streamWait(const L1& a,
|
|
176
|
-
|
|
177
|
-
streamWaitBase(a, b);
|
|
182
|
+
void streamWait(const L1& a, const std::initializer_list<cudaStream_t>& b) {
|
|
183
|
+
streamWaitBase(a, b);
|
|
178
184
|
}
|
|
179
185
|
|
|
180
186
|
template <typename L2>
|
|
181
|
-
void streamWait(const std::initializer_list<cudaStream_t>& a,
|
|
182
|
-
|
|
183
|
-
streamWaitBase(a, b);
|
|
187
|
+
void streamWait(const std::initializer_list<cudaStream_t>& a, const L2& b) {
|
|
188
|
+
streamWaitBase(a, b);
|
|
184
189
|
}
|
|
185
190
|
|
|
186
|
-
inline void streamWait(
|
|
187
|
-
|
|
188
|
-
|
|
191
|
+
inline void streamWait(
|
|
192
|
+
const std::initializer_list<cudaStream_t>& a,
|
|
193
|
+
const std::initializer_list<cudaStream_t>& b) {
|
|
194
|
+
streamWaitBase(a, b);
|
|
189
195
|
}
|
|
190
196
|
|
|
191
|
-
}
|
|
197
|
+
} // namespace gpu
|
|
198
|
+
} // namespace faiss
|