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,7 +5,6 @@
|
|
|
5
5
|
* LICENSE file in the root directory of this source tree.
|
|
6
6
|
*/
|
|
7
7
|
|
|
8
|
-
|
|
9
8
|
#pragma once
|
|
10
9
|
|
|
11
10
|
#include <faiss/gpu/GpuIndex.h>
|
|
@@ -17,189 +16,201 @@ struct IndexFlat;
|
|
|
17
16
|
struct IndexFlatL2;
|
|
18
17
|
struct IndexFlatIP;
|
|
19
18
|
|
|
20
|
-
}
|
|
19
|
+
} // namespace faiss
|
|
21
20
|
|
|
22
|
-
namespace faiss {
|
|
21
|
+
namespace faiss {
|
|
22
|
+
namespace gpu {
|
|
23
23
|
|
|
24
24
|
class FlatIndex;
|
|
25
25
|
|
|
26
26
|
struct GpuIndexFlatConfig : public GpuIndexConfig {
|
|
27
|
-
|
|
28
|
-
|
|
29
|
-
|
|
30
|
-
|
|
31
|
-
|
|
32
|
-
|
|
33
|
-
|
|
34
|
-
|
|
35
|
-
|
|
36
|
-
|
|
37
|
-
|
|
38
|
-
|
|
39
|
-
/// be transposed, and will increase storage requirements (we store
|
|
40
|
-
/// data in both transposed and non-transposed layouts).
|
|
41
|
-
bool storeTransposed;
|
|
27
|
+
inline GpuIndexFlatConfig() : useFloat16(false), storeTransposed(false) {}
|
|
28
|
+
|
|
29
|
+
/// Whether or not data is stored as float16
|
|
30
|
+
bool useFloat16;
|
|
31
|
+
|
|
32
|
+
/// Whether or not data is stored (transparently) in a transposed
|
|
33
|
+
/// layout, enabling use of the NN GEMM call, which is ~10% faster.
|
|
34
|
+
/// This will improve the speed of the flat index, but will
|
|
35
|
+
/// substantially slow down any add() calls made, as all data must
|
|
36
|
+
/// be transposed, and will increase storage requirements (we store
|
|
37
|
+
/// data in both transposed and non-transposed layouts).
|
|
38
|
+
bool storeTransposed;
|
|
42
39
|
};
|
|
43
40
|
|
|
44
41
|
/// Wrapper around the GPU implementation that looks like
|
|
45
42
|
/// faiss::IndexFlat; copies over centroid data from a given
|
|
46
43
|
/// faiss::IndexFlat
|
|
47
44
|
class GpuIndexFlat : public GpuIndex {
|
|
48
|
-
|
|
49
|
-
|
|
50
|
-
|
|
51
|
-
|
|
52
|
-
|
|
53
|
-
|
|
54
|
-
|
|
55
|
-
|
|
56
|
-
|
|
57
|
-
|
|
58
|
-
|
|
59
|
-
|
|
60
|
-
|
|
61
|
-
|
|
62
|
-
|
|
63
|
-
|
|
64
|
-
|
|
65
|
-
|
|
66
|
-
|
|
67
|
-
|
|
68
|
-
|
|
69
|
-
|
|
70
|
-
|
|
71
|
-
|
|
72
|
-
|
|
73
|
-
|
|
74
|
-
|
|
75
|
-
|
|
76
|
-
|
|
77
|
-
|
|
78
|
-
|
|
79
|
-
|
|
80
|
-
|
|
81
|
-
|
|
82
|
-
|
|
83
|
-
|
|
84
|
-
|
|
85
|
-
|
|
86
|
-
|
|
87
|
-
|
|
88
|
-
|
|
89
|
-
|
|
90
|
-
|
|
91
|
-
|
|
92
|
-
|
|
93
|
-
|
|
94
|
-
|
|
95
|
-
|
|
96
|
-
|
|
97
|
-
|
|
98
|
-
|
|
99
|
-
|
|
100
|
-
|
|
101
|
-
|
|
102
|
-
|
|
103
|
-
|
|
104
|
-
|
|
105
|
-
|
|
106
|
-
|
|
107
|
-
|
|
108
|
-
|
|
109
|
-
|
|
110
|
-
|
|
111
|
-
|
|
112
|
-
|
|
113
|
-
|
|
114
|
-
|
|
115
|
-
|
|
116
|
-
|
|
117
|
-
|
|
118
|
-
|
|
119
|
-
|
|
120
|
-
|
|
121
|
-
|
|
122
|
-
|
|
123
|
-
|
|
124
|
-
|
|
125
|
-
|
|
126
|
-
|
|
127
|
-
|
|
128
|
-
|
|
129
|
-
|
|
130
|
-
|
|
131
|
-
|
|
132
|
-
|
|
133
|
-
|
|
134
|
-
|
|
135
|
-
|
|
136
|
-
|
|
45
|
+
public:
|
|
46
|
+
/// Construct from a pre-existing faiss::IndexFlat instance, copying
|
|
47
|
+
/// data over to the given GPU
|
|
48
|
+
GpuIndexFlat(
|
|
49
|
+
GpuResourcesProvider* provider,
|
|
50
|
+
const faiss::IndexFlat* index,
|
|
51
|
+
GpuIndexFlatConfig config = GpuIndexFlatConfig());
|
|
52
|
+
|
|
53
|
+
GpuIndexFlat(
|
|
54
|
+
std::shared_ptr<GpuResources> resources,
|
|
55
|
+
const faiss::IndexFlat* index,
|
|
56
|
+
GpuIndexFlatConfig config = GpuIndexFlatConfig());
|
|
57
|
+
|
|
58
|
+
/// Construct an empty instance that can be added to
|
|
59
|
+
GpuIndexFlat(
|
|
60
|
+
GpuResourcesProvider* provider,
|
|
61
|
+
int dims,
|
|
62
|
+
faiss::MetricType metric,
|
|
63
|
+
GpuIndexFlatConfig config = GpuIndexFlatConfig());
|
|
64
|
+
|
|
65
|
+
GpuIndexFlat(
|
|
66
|
+
std::shared_ptr<GpuResources> resources,
|
|
67
|
+
int dims,
|
|
68
|
+
faiss::MetricType metric,
|
|
69
|
+
GpuIndexFlatConfig config = GpuIndexFlatConfig());
|
|
70
|
+
|
|
71
|
+
~GpuIndexFlat() override;
|
|
72
|
+
|
|
73
|
+
/// Initialize ourselves from the given CPU index; will overwrite
|
|
74
|
+
/// all data in ourselves
|
|
75
|
+
void copyFrom(const faiss::IndexFlat* index);
|
|
76
|
+
|
|
77
|
+
/// Copy ourselves to the given CPU index; will overwrite all data
|
|
78
|
+
/// in the index instance
|
|
79
|
+
void copyTo(faiss::IndexFlat* index) const;
|
|
80
|
+
|
|
81
|
+
/// Returns the number of vectors we contain
|
|
82
|
+
size_t getNumVecs() const;
|
|
83
|
+
|
|
84
|
+
/// Clears all vectors from this index
|
|
85
|
+
void reset() override;
|
|
86
|
+
|
|
87
|
+
/// This index is not trained, so this does nothing
|
|
88
|
+
void train(Index::idx_t n, const float* x) override;
|
|
89
|
+
|
|
90
|
+
/// Overrides to avoid excessive copies
|
|
91
|
+
void add(Index::idx_t, const float* x) override;
|
|
92
|
+
|
|
93
|
+
/// Reconstruction methods; prefer the batch reconstruct as it will
|
|
94
|
+
/// be more efficient
|
|
95
|
+
void reconstruct(Index::idx_t key, float* out) const override;
|
|
96
|
+
|
|
97
|
+
/// Batch reconstruction method
|
|
98
|
+
void reconstruct_n(Index::idx_t i0, Index::idx_t num, float* out)
|
|
99
|
+
const override;
|
|
100
|
+
|
|
101
|
+
/// Compute residual
|
|
102
|
+
void compute_residual(const float* x, float* residual, Index::idx_t key)
|
|
103
|
+
const override;
|
|
104
|
+
|
|
105
|
+
/// Compute residual (batch mode)
|
|
106
|
+
void compute_residual_n(
|
|
107
|
+
Index::idx_t n,
|
|
108
|
+
const float* xs,
|
|
109
|
+
float* residuals,
|
|
110
|
+
const Index::idx_t* keys) const override;
|
|
111
|
+
|
|
112
|
+
/// For internal access
|
|
113
|
+
inline FlatIndex* getGpuData() {
|
|
114
|
+
return data_.get();
|
|
115
|
+
}
|
|
116
|
+
|
|
117
|
+
protected:
|
|
118
|
+
/// Flat index does not require IDs as there is no storage available for
|
|
119
|
+
/// them
|
|
120
|
+
bool addImplRequiresIDs_() const override;
|
|
121
|
+
|
|
122
|
+
/// Called from GpuIndex for add
|
|
123
|
+
void addImpl_(int n, const float* x, const Index::idx_t* ids) override;
|
|
124
|
+
|
|
125
|
+
/// Called from GpuIndex for search
|
|
126
|
+
void searchImpl_(
|
|
127
|
+
int n,
|
|
128
|
+
const float* x,
|
|
129
|
+
int k,
|
|
130
|
+
float* distances,
|
|
131
|
+
Index::idx_t* labels) const override;
|
|
132
|
+
|
|
133
|
+
protected:
|
|
134
|
+
/// Our configuration options
|
|
135
|
+
const GpuIndexFlatConfig flatConfig_;
|
|
136
|
+
|
|
137
|
+
/// Holds our GPU data containing the list of vectors
|
|
138
|
+
std::unique_ptr<FlatIndex> data_;
|
|
137
139
|
};
|
|
138
140
|
|
|
139
141
|
/// Wrapper around the GPU implementation that looks like
|
|
140
142
|
/// faiss::IndexFlatL2; copies over centroid data from a given
|
|
141
143
|
/// faiss::IndexFlat
|
|
142
144
|
class GpuIndexFlatL2 : public GpuIndexFlat {
|
|
143
|
-
|
|
144
|
-
|
|
145
|
-
|
|
146
|
-
|
|
147
|
-
|
|
148
|
-
|
|
149
|
-
|
|
150
|
-
|
|
151
|
-
|
|
152
|
-
|
|
153
|
-
|
|
154
|
-
|
|
155
|
-
|
|
156
|
-
|
|
157
|
-
|
|
158
|
-
|
|
159
|
-
|
|
160
|
-
|
|
161
|
-
|
|
162
|
-
|
|
163
|
-
|
|
164
|
-
|
|
165
|
-
|
|
166
|
-
|
|
167
|
-
|
|
168
|
-
|
|
169
|
-
|
|
145
|
+
public:
|
|
146
|
+
/// Construct from a pre-existing faiss::IndexFlatL2 instance, copying
|
|
147
|
+
/// data over to the given GPU
|
|
148
|
+
GpuIndexFlatL2(
|
|
149
|
+
GpuResourcesProvider* provider,
|
|
150
|
+
faiss::IndexFlatL2* index,
|
|
151
|
+
GpuIndexFlatConfig config = GpuIndexFlatConfig());
|
|
152
|
+
|
|
153
|
+
GpuIndexFlatL2(
|
|
154
|
+
std::shared_ptr<GpuResources> resources,
|
|
155
|
+
faiss::IndexFlatL2* index,
|
|
156
|
+
GpuIndexFlatConfig config = GpuIndexFlatConfig());
|
|
157
|
+
|
|
158
|
+
/// Construct an empty instance that can be added to
|
|
159
|
+
GpuIndexFlatL2(
|
|
160
|
+
GpuResourcesProvider* provider,
|
|
161
|
+
int dims,
|
|
162
|
+
GpuIndexFlatConfig config = GpuIndexFlatConfig());
|
|
163
|
+
|
|
164
|
+
GpuIndexFlatL2(
|
|
165
|
+
std::shared_ptr<GpuResources> resources,
|
|
166
|
+
int dims,
|
|
167
|
+
GpuIndexFlatConfig config = GpuIndexFlatConfig());
|
|
168
|
+
|
|
169
|
+
/// Initialize ourselves from the given CPU index; will overwrite
|
|
170
|
+
/// all data in ourselves
|
|
171
|
+
void copyFrom(faiss::IndexFlat* index);
|
|
172
|
+
|
|
173
|
+
/// Copy ourselves to the given CPU index; will overwrite all data
|
|
174
|
+
/// in the index instance
|
|
175
|
+
void copyTo(faiss::IndexFlat* index);
|
|
170
176
|
};
|
|
171
177
|
|
|
172
178
|
/// Wrapper around the GPU implementation that looks like
|
|
173
179
|
/// faiss::IndexFlatIP; copies over centroid data from a given
|
|
174
180
|
/// faiss::IndexFlat
|
|
175
181
|
class GpuIndexFlatIP : public GpuIndexFlat {
|
|
176
|
-
|
|
177
|
-
|
|
178
|
-
|
|
179
|
-
|
|
180
|
-
|
|
181
|
-
|
|
182
|
-
|
|
183
|
-
|
|
184
|
-
|
|
185
|
-
|
|
186
|
-
|
|
187
|
-
|
|
188
|
-
|
|
189
|
-
|
|
190
|
-
|
|
191
|
-
|
|
192
|
-
|
|
193
|
-
|
|
194
|
-
|
|
195
|
-
|
|
196
|
-
|
|
197
|
-
|
|
198
|
-
|
|
199
|
-
|
|
200
|
-
|
|
201
|
-
|
|
202
|
-
|
|
182
|
+
public:
|
|
183
|
+
/// Construct from a pre-existing faiss::IndexFlatIP instance, copying
|
|
184
|
+
/// data over to the given GPU
|
|
185
|
+
GpuIndexFlatIP(
|
|
186
|
+
GpuResourcesProvider* provider,
|
|
187
|
+
faiss::IndexFlatIP* index,
|
|
188
|
+
GpuIndexFlatConfig config = GpuIndexFlatConfig());
|
|
189
|
+
|
|
190
|
+
GpuIndexFlatIP(
|
|
191
|
+
std::shared_ptr<GpuResources> resources,
|
|
192
|
+
faiss::IndexFlatIP* index,
|
|
193
|
+
GpuIndexFlatConfig config = GpuIndexFlatConfig());
|
|
194
|
+
|
|
195
|
+
/// Construct an empty instance that can be added to
|
|
196
|
+
GpuIndexFlatIP(
|
|
197
|
+
GpuResourcesProvider* provider,
|
|
198
|
+
int dims,
|
|
199
|
+
GpuIndexFlatConfig config = GpuIndexFlatConfig());
|
|
200
|
+
|
|
201
|
+
GpuIndexFlatIP(
|
|
202
|
+
std::shared_ptr<GpuResources> resources,
|
|
203
|
+
int dims,
|
|
204
|
+
GpuIndexFlatConfig config = GpuIndexFlatConfig());
|
|
205
|
+
|
|
206
|
+
/// Initialize ourselves from the given CPU index; will overwrite
|
|
207
|
+
/// all data in ourselves
|
|
208
|
+
void copyFrom(faiss::IndexFlat* index);
|
|
209
|
+
|
|
210
|
+
/// Copy ourselves to the given CPU index; will overwrite all data
|
|
211
|
+
/// in the index instance
|
|
212
|
+
void copyTo(faiss::IndexFlat* index);
|
|
203
213
|
};
|
|
204
214
|
|
|
205
|
-
}
|
|
215
|
+
} // namespace gpu
|
|
216
|
+
} // namespace faiss
|
|
@@ -5,102 +5,105 @@
|
|
|
5
5
|
* LICENSE file in the root directory of this source tree.
|
|
6
6
|
*/
|
|
7
7
|
|
|
8
|
-
|
|
9
8
|
#pragma once
|
|
10
9
|
|
|
10
|
+
#include <faiss/Clustering.h>
|
|
11
11
|
#include <faiss/gpu/GpuIndex.h>
|
|
12
12
|
#include <faiss/gpu/GpuIndexFlat.h>
|
|
13
13
|
#include <faiss/gpu/GpuIndicesOptions.h>
|
|
14
|
-
#include <faiss/Clustering.h>
|
|
15
14
|
|
|
16
|
-
namespace faiss {
|
|
15
|
+
namespace faiss {
|
|
16
|
+
struct IndexIVF;
|
|
17
|
+
}
|
|
17
18
|
|
|
18
|
-
namespace faiss {
|
|
19
|
+
namespace faiss {
|
|
20
|
+
namespace gpu {
|
|
19
21
|
|
|
20
22
|
class GpuIndexFlat;
|
|
21
23
|
|
|
22
24
|
struct GpuIndexIVFConfig : public GpuIndexConfig {
|
|
23
|
-
|
|
24
|
-
: indicesOptions(INDICES_64_BIT) {
|
|
25
|
-
}
|
|
25
|
+
inline GpuIndexIVFConfig() : indicesOptions(INDICES_64_BIT) {}
|
|
26
26
|
|
|
27
|
-
|
|
28
|
-
|
|
27
|
+
/// Index storage options for the GPU
|
|
28
|
+
IndicesOptions indicesOptions;
|
|
29
29
|
|
|
30
|
-
|
|
31
|
-
|
|
30
|
+
/// Configuration for the coarse quantizer object
|
|
31
|
+
GpuIndexFlatConfig flatConfig;
|
|
32
32
|
};
|
|
33
33
|
|
|
34
34
|
class GpuIndexIVF : public GpuIndex {
|
|
35
|
-
|
|
36
|
-
|
|
37
|
-
|
|
38
|
-
|
|
39
|
-
|
|
40
|
-
|
|
41
|
-
|
|
35
|
+
public:
|
|
36
|
+
GpuIndexIVF(
|
|
37
|
+
GpuResourcesProvider* provider,
|
|
38
|
+
int dims,
|
|
39
|
+
faiss::MetricType metric,
|
|
40
|
+
float metricArg,
|
|
41
|
+
int nlist,
|
|
42
|
+
GpuIndexIVFConfig config = GpuIndexIVFConfig());
|
|
42
43
|
|
|
43
|
-
|
|
44
|
+
~GpuIndexIVF() override;
|
|
44
45
|
|
|
45
|
-
|
|
46
|
-
|
|
47
|
-
|
|
46
|
+
private:
|
|
47
|
+
/// Shared initialization functions
|
|
48
|
+
void init_();
|
|
48
49
|
|
|
49
|
-
|
|
50
|
-
|
|
51
|
-
|
|
50
|
+
public:
|
|
51
|
+
/// Copy what we need from the CPU equivalent
|
|
52
|
+
void copyFrom(const faiss::IndexIVF* index);
|
|
52
53
|
|
|
53
|
-
|
|
54
|
-
|
|
54
|
+
/// Copy what we have to the CPU equivalent
|
|
55
|
+
void copyTo(faiss::IndexIVF* index) const;
|
|
55
56
|
|
|
56
|
-
|
|
57
|
-
|
|
57
|
+
/// Returns the number of inverted lists we're managing
|
|
58
|
+
int getNumLists() const;
|
|
58
59
|
|
|
59
|
-
|
|
60
|
-
|
|
60
|
+
/// Returns the number of vectors present in a particular inverted list
|
|
61
|
+
virtual int getListLength(int listId) const = 0;
|
|
61
62
|
|
|
62
|
-
|
|
63
|
-
|
|
64
|
-
|
|
65
|
-
|
|
66
|
-
|
|
67
|
-
|
|
68
|
-
|
|
69
|
-
|
|
63
|
+
/// Return the encoded vector data contained in a particular inverted list,
|
|
64
|
+
/// for debugging purposes.
|
|
65
|
+
/// If gpuFormat is true, the data is returned as it is encoded in the
|
|
66
|
+
/// GPU-side representation.
|
|
67
|
+
/// Otherwise, it is converted to the CPU format.
|
|
68
|
+
/// compliant format, while the native GPU format may differ.
|
|
69
|
+
virtual std::vector<uint8_t> getListVectorData(
|
|
70
|
+
int listId,
|
|
71
|
+
bool gpuFormat = false) const = 0;
|
|
70
72
|
|
|
71
|
-
|
|
72
|
-
|
|
73
|
-
|
|
73
|
+
/// Return the vector indices contained in a particular inverted list, for
|
|
74
|
+
/// debugging purposes.
|
|
75
|
+
virtual std::vector<Index::idx_t> getListIndices(int listId) const = 0;
|
|
74
76
|
|
|
75
|
-
|
|
76
|
-
|
|
77
|
+
/// Return the quantizer we're using
|
|
78
|
+
GpuIndexFlat* getQuantizer();
|
|
77
79
|
|
|
78
|
-
|
|
79
|
-
|
|
80
|
+
/// Sets the number of list probes per query
|
|
81
|
+
void setNumProbes(int nprobe);
|
|
80
82
|
|
|
81
|
-
|
|
82
|
-
|
|
83
|
+
/// Returns our current number of list probes per query
|
|
84
|
+
int getNumProbes() const;
|
|
83
85
|
|
|
84
|
-
|
|
85
|
-
|
|
86
|
-
|
|
86
|
+
protected:
|
|
87
|
+
bool addImplRequiresIDs_() const override;
|
|
88
|
+
void trainQuantizer_(Index::idx_t n, const float* x);
|
|
87
89
|
|
|
88
|
-
|
|
89
|
-
|
|
90
|
-
|
|
90
|
+
public:
|
|
91
|
+
/// Exposing this like the CPU version for manipulation
|
|
92
|
+
ClusteringParameters cp;
|
|
91
93
|
|
|
92
|
-
|
|
93
|
-
|
|
94
|
+
/// Exposing this like the CPU version for query
|
|
95
|
+
int nlist;
|
|
94
96
|
|
|
95
|
-
|
|
96
|
-
|
|
97
|
+
/// Exposing this like the CPU version for manipulation
|
|
98
|
+
int nprobe;
|
|
97
99
|
|
|
98
|
-
|
|
99
|
-
|
|
100
|
+
/// Exposeing this like the CPU version for query
|
|
101
|
+
GpuIndexFlat* quantizer;
|
|
100
102
|
|
|
101
|
-
|
|
102
|
-
|
|
103
|
-
|
|
103
|
+
protected:
|
|
104
|
+
/// Our configuration options
|
|
105
|
+
const GpuIndexIVFConfig ivfConfig_;
|
|
104
106
|
};
|
|
105
107
|
|
|
106
|
-
}
|
|
108
|
+
} // namespace gpu
|
|
109
|
+
} // namespace faiss
|