faiss 0.1.5 → 0.2.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 +24 -0
- data/README.md +12 -0
- data/ext/faiss/ext.cpp +1 -1
- data/ext/faiss/extconf.rb +6 -2
- data/ext/faiss/index.cpp +114 -43
- data/ext/faiss/index_binary.cpp +24 -30
- data/ext/faiss/kmeans.cpp +20 -16
- data/ext/faiss/numo.hpp +867 -0
- data/ext/faiss/pca_matrix.cpp +13 -14
- data/ext/faiss/product_quantizer.cpp +23 -24
- data/ext/faiss/utils.cpp +10 -37
- data/ext/faiss/utils.h +2 -13
- data/lib/faiss.rb +0 -5
- 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 +334 -195
- data/vendor/faiss/faiss/Clustering.h +88 -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 +139 -164
- data/vendor/faiss/faiss/Index2Layer.h +22 -22
- 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 +116 -147
- data/vendor/faiss/faiss/IndexFlat.h +35 -46
- data/vendor/faiss/faiss/IndexHNSW.cpp +372 -348
- data/vendor/faiss/faiss/IndexHNSW.h +57 -41
- data/vendor/faiss/faiss/IndexIVF.cpp +474 -454
- data/vendor/faiss/faiss/IndexIVF.h +146 -113
- data/vendor/faiss/faiss/IndexIVFFlat.cpp +248 -250
- data/vendor/faiss/faiss/IndexIVFFlat.h +48 -51
- data/vendor/faiss/faiss/IndexIVFPQ.cpp +457 -516
- data/vendor/faiss/faiss/IndexIVFPQ.h +74 -66
- 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 +125 -133
- data/vendor/faiss/faiss/IndexIVFSpectralHash.h +19 -21
- data/vendor/faiss/faiss/IndexLSH.cpp +75 -96
- data/vendor/faiss/faiss/IndexLSH.h +21 -26
- data/vendor/faiss/faiss/IndexLattice.cpp +42 -56
- data/vendor/faiss/faiss/IndexLattice.h +11 -16
- data/vendor/faiss/faiss/IndexNNDescent.cpp +231 -0
- data/vendor/faiss/faiss/IndexNNDescent.h +72 -0
- data/vendor/faiss/faiss/IndexNSG.cpp +303 -0
- data/vendor/faiss/faiss/IndexNSG.h +85 -0
- data/vendor/faiss/faiss/IndexPQ.cpp +405 -464
- data/vendor/faiss/faiss/IndexPQ.h +64 -67
- 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 +115 -131
- data/vendor/faiss/faiss/IndexRefine.h +22 -23
- data/vendor/faiss/faiss/IndexReplicas.cpp +147 -153
- data/vendor/faiss/faiss/IndexReplicas.h +62 -56
- data/vendor/faiss/faiss/IndexResidual.cpp +291 -0
- data/vendor/faiss/faiss/IndexResidual.h +152 -0
- data/vendor/faiss/faiss/IndexScalarQuantizer.cpp +120 -155
- data/vendor/faiss/faiss/IndexScalarQuantizer.h +41 -45
- 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 +652 -474
- data/vendor/faiss/faiss/VectorTransform.h +61 -89
- data/vendor/faiss/faiss/clone_index.cpp +77 -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 +197 -170
- 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/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 +270 -0
- data/vendor/faiss/faiss/impl/AdditiveQuantizer.h +115 -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 +595 -611
- data/vendor/faiss/faiss/impl/HNSW.h +179 -200
- data/vendor/faiss/faiss/impl/LocalSearchQuantizer.cpp +672 -0
- data/vendor/faiss/faiss/impl/LocalSearchQuantizer.h +172 -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 +682 -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 +448 -0
- data/vendor/faiss/faiss/impl/ResidualQuantizer.h +130 -0
- data/vendor/faiss/faiss/impl/ResultHandler.h +96 -132
- data/vendor/faiss/faiss/impl/ScalarQuantizer.cpp +648 -701
- 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 +547 -479
- data/vendor/faiss/faiss/impl/index_write.cpp +497 -407
- data/vendor/faiss/faiss/impl/io.cpp +75 -94
- data/vendor/faiss/faiss/impl/io.h +31 -41
- data/vendor/faiss/faiss/impl/io_macros.h +40 -29
- 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 +269 -218
- data/vendor/faiss/faiss/index_factory.h +6 -7
- 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 +301 -310
- data/vendor/faiss/faiss/utils/distances.h +133 -118
- data/vendor/faiss/faiss/utils/distances_simd.cpp +456 -516
- 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 +53 -48
- metadata +24 -10
- data/lib/faiss/index.rb +0 -20
- data/lib/faiss/index_binary.rb +0 -20
- data/lib/faiss/kmeans.rb +0 -15
- data/lib/faiss/pca_matrix.rb +0 -15
- data/lib/faiss/product_quantizer.rb +0 -22
@@ -5,12 +5,11 @@
|
|
5
5
|
* LICENSE file in the root directory of this source tree.
|
6
6
|
*/
|
7
7
|
|
8
|
-
|
9
8
|
#ifndef GPU_FAISS_ASSERT_INCLUDED
|
10
9
|
#define GPU_FAISS_ASSERT_INCLUDED
|
11
10
|
|
12
|
-
#include <faiss/impl/FaissAssert.h>
|
13
11
|
#include <cuda.h>
|
12
|
+
#include <faiss/impl/FaissAssert.h>
|
14
13
|
|
15
14
|
///
|
16
15
|
/// Assertions
|
@@ -5,152 +5,149 @@
|
|
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/Index.h>
|
12
11
|
#include <faiss/gpu/GpuResources.h>
|
13
12
|
|
14
|
-
namespace faiss {
|
13
|
+
namespace faiss {
|
14
|
+
namespace gpu {
|
15
15
|
|
16
16
|
struct GpuIndexConfig {
|
17
|
-
|
18
|
-
|
19
|
-
|
20
|
-
|
21
|
-
|
22
|
-
|
23
|
-
|
24
|
-
|
25
|
-
|
26
|
-
/// On Pascal and above (CC 6+) architectures, allows GPUs to use
|
27
|
-
/// more memory than is available on the GPU.
|
28
|
-
MemorySpace memorySpace;
|
17
|
+
inline GpuIndexConfig() : device(0), memorySpace(MemorySpace::Device) {}
|
18
|
+
|
19
|
+
/// GPU device on which the index is resident
|
20
|
+
int device;
|
21
|
+
|
22
|
+
/// What memory space to use for primary storage.
|
23
|
+
/// On Pascal and above (CC 6+) architectures, allows GPUs to use
|
24
|
+
/// more memory than is available on the GPU.
|
25
|
+
MemorySpace memorySpace;
|
29
26
|
};
|
30
27
|
|
31
28
|
class GpuIndex : public faiss::Index {
|
32
|
-
|
33
|
-
|
34
|
-
|
35
|
-
|
36
|
-
|
37
|
-
|
38
|
-
|
39
|
-
|
40
|
-
|
41
|
-
|
42
|
-
|
43
|
-
|
44
|
-
|
45
|
-
|
46
|
-
|
47
|
-
|
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
|
-
|
137
|
-
|
138
|
-
|
139
|
-
|
140
|
-
|
141
|
-
|
142
|
-
|
143
|
-
|
144
|
-
|
145
|
-
|
146
|
-
|
147
|
-
|
148
|
-
|
149
|
-
|
150
|
-
|
151
|
-
|
152
|
-
|
153
|
-
size_t minPagedSize_;
|
29
|
+
public:
|
30
|
+
GpuIndex(
|
31
|
+
std::shared_ptr<GpuResources> resources,
|
32
|
+
int dims,
|
33
|
+
faiss::MetricType metric,
|
34
|
+
float metricArg,
|
35
|
+
GpuIndexConfig config);
|
36
|
+
|
37
|
+
/// Returns the device that this index is resident on
|
38
|
+
int getDevice() const;
|
39
|
+
|
40
|
+
/// Returns a reference to our GpuResources object that manages memory,
|
41
|
+
/// stream and handle resources on the GPU
|
42
|
+
std::shared_ptr<GpuResources> getResources();
|
43
|
+
|
44
|
+
/// Set the minimum data size for searches (in MiB) for which we use
|
45
|
+
/// CPU -> GPU paging
|
46
|
+
void setMinPagingSize(size_t size);
|
47
|
+
|
48
|
+
/// Returns the current minimum data size for paged searches
|
49
|
+
size_t getMinPagingSize() const;
|
50
|
+
|
51
|
+
/// `x` can be resident on the CPU or any GPU; copies are performed
|
52
|
+
/// as needed
|
53
|
+
/// Handles paged adds if the add set is too large; calls addInternal_
|
54
|
+
void add(Index::idx_t, const float* x) override;
|
55
|
+
|
56
|
+
/// `x` and `ids` can be resident on the CPU or any GPU; copies are
|
57
|
+
/// performed as needed
|
58
|
+
/// Handles paged adds if the add set is too large; calls addInternal_
|
59
|
+
void add_with_ids(Index::idx_t n, const float* x, const Index::idx_t* ids)
|
60
|
+
override;
|
61
|
+
|
62
|
+
/// `x` and `labels` can be resident on the CPU or any GPU; copies are
|
63
|
+
/// performed as needed
|
64
|
+
void assign(
|
65
|
+
Index::idx_t n,
|
66
|
+
const float* x,
|
67
|
+
Index::idx_t* labels,
|
68
|
+
Index::idx_t k = 1) const override;
|
69
|
+
|
70
|
+
/// `x`, `distances` and `labels` can be resident on the CPU or any
|
71
|
+
/// GPU; copies are performed as needed
|
72
|
+
void search(
|
73
|
+
Index::idx_t n,
|
74
|
+
const float* x,
|
75
|
+
Index::idx_t k,
|
76
|
+
float* distances,
|
77
|
+
Index::idx_t* labels) const override;
|
78
|
+
|
79
|
+
/// Overridden to force GPU indices to provide their own GPU-friendly
|
80
|
+
/// implementation
|
81
|
+
void compute_residual(const float* x, float* residual, Index::idx_t key)
|
82
|
+
const override;
|
83
|
+
|
84
|
+
/// Overridden to force GPU indices to provide their own GPU-friendly
|
85
|
+
/// implementation
|
86
|
+
void compute_residual_n(
|
87
|
+
Index::idx_t n,
|
88
|
+
const float* xs,
|
89
|
+
float* residuals,
|
90
|
+
const Index::idx_t* keys) const override;
|
91
|
+
|
92
|
+
protected:
|
93
|
+
/// Copy what we need from the CPU equivalent
|
94
|
+
void copyFrom(const faiss::Index* index);
|
95
|
+
|
96
|
+
/// Copy what we have to the CPU equivalent
|
97
|
+
void copyTo(faiss::Index* index) const;
|
98
|
+
|
99
|
+
/// Does addImpl_ require IDs? If so, and no IDs are provided, we will
|
100
|
+
/// generate them sequentially based on the order in which the IDs are added
|
101
|
+
virtual bool addImplRequiresIDs_() const = 0;
|
102
|
+
|
103
|
+
/// Overridden to actually perform the add
|
104
|
+
/// All data is guaranteed to be resident on our device
|
105
|
+
virtual void addImpl_(int n, const float* x, const Index::idx_t* ids) = 0;
|
106
|
+
|
107
|
+
/// Overridden to actually perform the search
|
108
|
+
/// All data is guaranteed to be resident on our device
|
109
|
+
virtual void searchImpl_(
|
110
|
+
int n,
|
111
|
+
const float* x,
|
112
|
+
int k,
|
113
|
+
float* distances,
|
114
|
+
Index::idx_t* labels) const = 0;
|
115
|
+
|
116
|
+
private:
|
117
|
+
/// Handles paged adds if the add set is too large, passes to
|
118
|
+
/// addImpl_ to actually perform the add for the current page
|
119
|
+
void addPaged_(int n, const float* x, const Index::idx_t* ids);
|
120
|
+
|
121
|
+
/// Calls addImpl_ for a single page of GPU-resident data
|
122
|
+
void addPage_(int n, const float* x, const Index::idx_t* ids);
|
123
|
+
|
124
|
+
/// Calls searchImpl_ for a single page of GPU-resident data
|
125
|
+
void searchNonPaged_(
|
126
|
+
int n,
|
127
|
+
const float* x,
|
128
|
+
int k,
|
129
|
+
float* outDistancesData,
|
130
|
+
Index::idx_t* outIndicesData) const;
|
131
|
+
|
132
|
+
/// Calls searchImpl_ for a single page of GPU-resident data,
|
133
|
+
/// handling paging of the data and copies from the CPU
|
134
|
+
void searchFromCpuPaged_(
|
135
|
+
int n,
|
136
|
+
const float* x,
|
137
|
+
int k,
|
138
|
+
float* outDistancesData,
|
139
|
+
Index::idx_t* outIndicesData) const;
|
140
|
+
|
141
|
+
protected:
|
142
|
+
/// Manages streams, cuBLAS handles and scratch memory for devices
|
143
|
+
std::shared_ptr<GpuResources> resources_;
|
144
|
+
|
145
|
+
/// Our configuration options
|
146
|
+
const GpuIndexConfig config_;
|
147
|
+
|
148
|
+
/// Size above which we page copies from the CPU to GPU
|
149
|
+
size_t minPagedSize_;
|
154
150
|
};
|
155
151
|
|
156
|
-
}
|
152
|
+
} // namespace gpu
|
153
|
+
} // namespace faiss
|
@@ -12,85 +12,88 @@
|
|
12
12
|
#include <faiss/gpu/GpuResources.h>
|
13
13
|
#include <memory>
|
14
14
|
|
15
|
-
namespace faiss {
|
15
|
+
namespace faiss {
|
16
|
+
namespace gpu {
|
16
17
|
|
17
18
|
class BinaryFlatIndex;
|
18
19
|
|
19
|
-
struct GpuIndexBinaryFlatConfig : public GpuIndexConfig {
|
20
|
-
};
|
20
|
+
struct GpuIndexBinaryFlatConfig : public GpuIndexConfig {};
|
21
21
|
|
22
22
|
/// A GPU version of IndexBinaryFlat for brute-force comparison of bit vectors
|
23
23
|
/// via Hamming distance
|
24
24
|
class GpuIndexBinaryFlat : public IndexBinary {
|
25
|
-
|
26
|
-
|
27
|
-
|
28
|
-
|
29
|
-
|
30
|
-
|
31
|
-
|
32
|
-
|
33
|
-
|
34
|
-
|
35
|
-
|
36
|
-
|
37
|
-
|
38
|
-
|
39
|
-
|
40
|
-
|
41
|
-
|
42
|
-
|
43
|
-
|
44
|
-
|
45
|
-
|
46
|
-
|
47
|
-
|
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
|
-
|
25
|
+
public:
|
26
|
+
/// Construct from a pre-existing faiss::IndexBinaryFlat instance, copying
|
27
|
+
/// data over to the given GPU
|
28
|
+
GpuIndexBinaryFlat(
|
29
|
+
GpuResourcesProvider* resources,
|
30
|
+
const faiss::IndexBinaryFlat* index,
|
31
|
+
GpuIndexBinaryFlatConfig config = GpuIndexBinaryFlatConfig());
|
32
|
+
|
33
|
+
/// Construct an empty instance that can be added to
|
34
|
+
GpuIndexBinaryFlat(
|
35
|
+
GpuResourcesProvider* resources,
|
36
|
+
int dims,
|
37
|
+
GpuIndexBinaryFlatConfig config = GpuIndexBinaryFlatConfig());
|
38
|
+
|
39
|
+
~GpuIndexBinaryFlat() override;
|
40
|
+
|
41
|
+
/// Returns the device that this index is resident on
|
42
|
+
int getDevice() const;
|
43
|
+
|
44
|
+
/// Returns a reference to our GpuResources object that manages memory,
|
45
|
+
/// stream and handle resources on the GPU
|
46
|
+
std::shared_ptr<GpuResources> getResources();
|
47
|
+
|
48
|
+
/// Initialize ourselves from the given CPU index; will overwrite
|
49
|
+
/// all data in ourselves
|
50
|
+
void copyFrom(const faiss::IndexBinaryFlat* index);
|
51
|
+
|
52
|
+
/// Copy ourselves to the given CPU index; will overwrite all data
|
53
|
+
/// in the index instance
|
54
|
+
void copyTo(faiss::IndexBinaryFlat* index) const;
|
55
|
+
|
56
|
+
void add(faiss::IndexBinary::idx_t n, const uint8_t* x) override;
|
57
|
+
|
58
|
+
void reset() override;
|
59
|
+
|
60
|
+
void search(
|
61
|
+
faiss::IndexBinary::idx_t n,
|
62
|
+
const uint8_t* x,
|
63
|
+
faiss::IndexBinary::idx_t k,
|
64
|
+
int32_t* distances,
|
65
|
+
faiss::IndexBinary::idx_t* labels) const override;
|
66
|
+
|
67
|
+
void reconstruct(faiss::IndexBinary::idx_t key, uint8_t* recons)
|
68
|
+
const override;
|
69
|
+
|
70
|
+
protected:
|
71
|
+
/// Called from search when the input data is on the CPU;
|
72
|
+
/// potentially allows for pinned memory usage
|
73
|
+
void searchFromCpuPaged_(
|
74
|
+
int n,
|
75
|
+
const uint8_t* x,
|
76
|
+
int k,
|
77
|
+
int32_t* outDistancesData,
|
78
|
+
int* outIndicesData) const;
|
79
|
+
|
80
|
+
void searchNonPaged_(
|
81
|
+
int n,
|
82
|
+
const uint8_t* x,
|
83
|
+
int k,
|
84
|
+
int32_t* outDistancesData,
|
85
|
+
int* outIndicesData) const;
|
86
|
+
|
87
|
+
protected:
|
88
|
+
/// Manages streans, cuBLAS handles and scratch memory for devices
|
89
|
+
std::shared_ptr<GpuResources> resources_;
|
90
|
+
|
91
|
+
/// Configuration options
|
92
|
+
const GpuIndexBinaryFlatConfig binaryFlatConfig_;
|
93
|
+
|
94
|
+
/// Holds our GPU data containing the list of vectors
|
95
|
+
std::unique_ptr<BinaryFlatIndex> data_;
|
94
96
|
};
|
95
97
|
|
96
|
-
}
|
98
|
+
} // namespace gpu
|
99
|
+
} // namespace faiss
|