faiss 0.2.0 → 0.2.1

Sign up to get free protection for your applications and to get access to all the features.
Files changed (202) hide show
  1. checksums.yaml +4 -4
  2. data/CHANGELOG.md +4 -0
  3. data/lib/faiss/version.rb +1 -1
  4. data/vendor/faiss/faiss/AutoTune.cpp +292 -291
  5. data/vendor/faiss/faiss/AutoTune.h +55 -56
  6. data/vendor/faiss/faiss/Clustering.cpp +334 -195
  7. data/vendor/faiss/faiss/Clustering.h +88 -35
  8. data/vendor/faiss/faiss/IVFlib.cpp +171 -195
  9. data/vendor/faiss/faiss/IVFlib.h +48 -51
  10. data/vendor/faiss/faiss/Index.cpp +85 -103
  11. data/vendor/faiss/faiss/Index.h +54 -48
  12. data/vendor/faiss/faiss/Index2Layer.cpp +139 -164
  13. data/vendor/faiss/faiss/Index2Layer.h +22 -22
  14. data/vendor/faiss/faiss/IndexBinary.cpp +45 -37
  15. data/vendor/faiss/faiss/IndexBinary.h +140 -132
  16. data/vendor/faiss/faiss/IndexBinaryFlat.cpp +73 -53
  17. data/vendor/faiss/faiss/IndexBinaryFlat.h +29 -24
  18. data/vendor/faiss/faiss/IndexBinaryFromFloat.cpp +46 -43
  19. data/vendor/faiss/faiss/IndexBinaryFromFloat.h +16 -15
  20. data/vendor/faiss/faiss/IndexBinaryHNSW.cpp +215 -232
  21. data/vendor/faiss/faiss/IndexBinaryHNSW.h +25 -24
  22. data/vendor/faiss/faiss/IndexBinaryHash.cpp +182 -177
  23. data/vendor/faiss/faiss/IndexBinaryHash.h +41 -34
  24. data/vendor/faiss/faiss/IndexBinaryIVF.cpp +489 -461
  25. data/vendor/faiss/faiss/IndexBinaryIVF.h +97 -68
  26. data/vendor/faiss/faiss/IndexFlat.cpp +116 -147
  27. data/vendor/faiss/faiss/IndexFlat.h +35 -46
  28. data/vendor/faiss/faiss/IndexHNSW.cpp +372 -348
  29. data/vendor/faiss/faiss/IndexHNSW.h +57 -41
  30. data/vendor/faiss/faiss/IndexIVF.cpp +474 -454
  31. data/vendor/faiss/faiss/IndexIVF.h +146 -113
  32. data/vendor/faiss/faiss/IndexIVFFlat.cpp +248 -250
  33. data/vendor/faiss/faiss/IndexIVFFlat.h +48 -51
  34. data/vendor/faiss/faiss/IndexIVFPQ.cpp +457 -516
  35. data/vendor/faiss/faiss/IndexIVFPQ.h +74 -66
  36. data/vendor/faiss/faiss/IndexIVFPQFastScan.cpp +406 -372
  37. data/vendor/faiss/faiss/IndexIVFPQFastScan.h +82 -57
  38. data/vendor/faiss/faiss/IndexIVFPQR.cpp +104 -102
  39. data/vendor/faiss/faiss/IndexIVFPQR.h +33 -28
  40. data/vendor/faiss/faiss/IndexIVFSpectralHash.cpp +125 -133
  41. data/vendor/faiss/faiss/IndexIVFSpectralHash.h +19 -21
  42. data/vendor/faiss/faiss/IndexLSH.cpp +75 -96
  43. data/vendor/faiss/faiss/IndexLSH.h +21 -26
  44. data/vendor/faiss/faiss/IndexLattice.cpp +42 -56
  45. data/vendor/faiss/faiss/IndexLattice.h +11 -16
  46. data/vendor/faiss/faiss/IndexNNDescent.cpp +231 -0
  47. data/vendor/faiss/faiss/IndexNNDescent.h +72 -0
  48. data/vendor/faiss/faiss/IndexNSG.cpp +303 -0
  49. data/vendor/faiss/faiss/IndexNSG.h +85 -0
  50. data/vendor/faiss/faiss/IndexPQ.cpp +405 -464
  51. data/vendor/faiss/faiss/IndexPQ.h +64 -67
  52. data/vendor/faiss/faiss/IndexPQFastScan.cpp +143 -170
  53. data/vendor/faiss/faiss/IndexPQFastScan.h +46 -32
  54. data/vendor/faiss/faiss/IndexPreTransform.cpp +120 -150
  55. data/vendor/faiss/faiss/IndexPreTransform.h +33 -36
  56. data/vendor/faiss/faiss/IndexRefine.cpp +115 -131
  57. data/vendor/faiss/faiss/IndexRefine.h +22 -23
  58. data/vendor/faiss/faiss/IndexReplicas.cpp +147 -153
  59. data/vendor/faiss/faiss/IndexReplicas.h +62 -56
  60. data/vendor/faiss/faiss/IndexResidual.cpp +291 -0
  61. data/vendor/faiss/faiss/IndexResidual.h +152 -0
  62. data/vendor/faiss/faiss/IndexScalarQuantizer.cpp +120 -155
  63. data/vendor/faiss/faiss/IndexScalarQuantizer.h +41 -45
  64. data/vendor/faiss/faiss/IndexShards.cpp +256 -240
  65. data/vendor/faiss/faiss/IndexShards.h +85 -73
  66. data/vendor/faiss/faiss/MatrixStats.cpp +112 -97
  67. data/vendor/faiss/faiss/MatrixStats.h +7 -10
  68. data/vendor/faiss/faiss/MetaIndexes.cpp +135 -157
  69. data/vendor/faiss/faiss/MetaIndexes.h +40 -34
  70. data/vendor/faiss/faiss/MetricType.h +7 -7
  71. data/vendor/faiss/faiss/VectorTransform.cpp +652 -474
  72. data/vendor/faiss/faiss/VectorTransform.h +61 -89
  73. data/vendor/faiss/faiss/clone_index.cpp +77 -73
  74. data/vendor/faiss/faiss/clone_index.h +4 -9
  75. data/vendor/faiss/faiss/gpu/GpuAutoTune.cpp +33 -38
  76. data/vendor/faiss/faiss/gpu/GpuAutoTune.h +11 -9
  77. data/vendor/faiss/faiss/gpu/GpuCloner.cpp +197 -170
  78. data/vendor/faiss/faiss/gpu/GpuCloner.h +53 -35
  79. data/vendor/faiss/faiss/gpu/GpuClonerOptions.cpp +12 -14
  80. data/vendor/faiss/faiss/gpu/GpuClonerOptions.h +27 -25
  81. data/vendor/faiss/faiss/gpu/GpuDistance.h +116 -112
  82. data/vendor/faiss/faiss/gpu/GpuFaissAssert.h +1 -2
  83. data/vendor/faiss/faiss/gpu/GpuIndex.h +134 -137
  84. data/vendor/faiss/faiss/gpu/GpuIndexBinaryFlat.h +76 -73
  85. data/vendor/faiss/faiss/gpu/GpuIndexFlat.h +173 -162
  86. data/vendor/faiss/faiss/gpu/GpuIndexIVF.h +67 -64
  87. data/vendor/faiss/faiss/gpu/GpuIndexIVFFlat.h +89 -86
  88. data/vendor/faiss/faiss/gpu/GpuIndexIVFPQ.h +150 -141
  89. data/vendor/faiss/faiss/gpu/GpuIndexIVFScalarQuantizer.h +101 -103
  90. data/vendor/faiss/faiss/gpu/GpuIndicesOptions.h +17 -16
  91. data/vendor/faiss/faiss/gpu/GpuResources.cpp +116 -128
  92. data/vendor/faiss/faiss/gpu/GpuResources.h +182 -186
  93. data/vendor/faiss/faiss/gpu/StandardGpuResources.cpp +433 -422
  94. data/vendor/faiss/faiss/gpu/StandardGpuResources.h +131 -130
  95. data/vendor/faiss/faiss/gpu/impl/InterleavedCodes.cpp +468 -456
  96. data/vendor/faiss/faiss/gpu/impl/InterleavedCodes.h +25 -19
  97. data/vendor/faiss/faiss/gpu/impl/RemapIndices.cpp +22 -20
  98. data/vendor/faiss/faiss/gpu/impl/RemapIndices.h +9 -8
  99. data/vendor/faiss/faiss/gpu/perf/IndexWrapper-inl.h +39 -44
  100. data/vendor/faiss/faiss/gpu/perf/IndexWrapper.h +16 -14
  101. data/vendor/faiss/faiss/gpu/perf/PerfClustering.cpp +77 -71
  102. data/vendor/faiss/faiss/gpu/perf/PerfIVFPQAdd.cpp +109 -88
  103. data/vendor/faiss/faiss/gpu/perf/WriteIndex.cpp +75 -64
  104. data/vendor/faiss/faiss/gpu/test/TestCodePacking.cpp +230 -215
  105. data/vendor/faiss/faiss/gpu/test/TestGpuIndexBinaryFlat.cpp +80 -86
  106. data/vendor/faiss/faiss/gpu/test/TestGpuIndexFlat.cpp +284 -277
  107. data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFFlat.cpp +416 -416
  108. data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFPQ.cpp +611 -517
  109. data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFScalarQuantizer.cpp +166 -164
  110. data/vendor/faiss/faiss/gpu/test/TestGpuMemoryException.cpp +61 -53
  111. data/vendor/faiss/faiss/gpu/test/TestUtils.cpp +274 -238
  112. data/vendor/faiss/faiss/gpu/test/TestUtils.h +73 -57
  113. data/vendor/faiss/faiss/gpu/test/demo_ivfpq_indexing_gpu.cpp +47 -50
  114. data/vendor/faiss/faiss/gpu/utils/DeviceUtils.h +79 -72
  115. data/vendor/faiss/faiss/gpu/utils/StackDeviceMemory.cpp +140 -146
  116. data/vendor/faiss/faiss/gpu/utils/StackDeviceMemory.h +69 -71
  117. data/vendor/faiss/faiss/gpu/utils/StaticUtils.h +21 -16
  118. data/vendor/faiss/faiss/gpu/utils/Timer.cpp +25 -29
  119. data/vendor/faiss/faiss/gpu/utils/Timer.h +30 -29
  120. data/vendor/faiss/faiss/impl/AdditiveQuantizer.cpp +270 -0
  121. data/vendor/faiss/faiss/impl/AdditiveQuantizer.h +115 -0
  122. data/vendor/faiss/faiss/impl/AuxIndexStructures.cpp +90 -120
  123. data/vendor/faiss/faiss/impl/AuxIndexStructures.h +81 -65
  124. data/vendor/faiss/faiss/impl/FaissAssert.h +73 -58
  125. data/vendor/faiss/faiss/impl/FaissException.cpp +56 -48
  126. data/vendor/faiss/faiss/impl/FaissException.h +41 -29
  127. data/vendor/faiss/faiss/impl/HNSW.cpp +595 -611
  128. data/vendor/faiss/faiss/impl/HNSW.h +179 -200
  129. data/vendor/faiss/faiss/impl/LocalSearchQuantizer.cpp +672 -0
  130. data/vendor/faiss/faiss/impl/LocalSearchQuantizer.h +172 -0
  131. data/vendor/faiss/faiss/impl/NNDescent.cpp +487 -0
  132. data/vendor/faiss/faiss/impl/NNDescent.h +154 -0
  133. data/vendor/faiss/faiss/impl/NSG.cpp +682 -0
  134. data/vendor/faiss/faiss/impl/NSG.h +199 -0
  135. data/vendor/faiss/faiss/impl/PolysemousTraining.cpp +484 -454
  136. data/vendor/faiss/faiss/impl/PolysemousTraining.h +52 -55
  137. data/vendor/faiss/faiss/impl/ProductQuantizer-inl.h +26 -47
  138. data/vendor/faiss/faiss/impl/ProductQuantizer.cpp +469 -459
  139. data/vendor/faiss/faiss/impl/ProductQuantizer.h +76 -87
  140. data/vendor/faiss/faiss/impl/ResidualQuantizer.cpp +448 -0
  141. data/vendor/faiss/faiss/impl/ResidualQuantizer.h +130 -0
  142. data/vendor/faiss/faiss/impl/ResultHandler.h +96 -132
  143. data/vendor/faiss/faiss/impl/ScalarQuantizer.cpp +648 -701
  144. data/vendor/faiss/faiss/impl/ScalarQuantizer.h +48 -46
  145. data/vendor/faiss/faiss/impl/ThreadedIndex-inl.h +129 -131
  146. data/vendor/faiss/faiss/impl/ThreadedIndex.h +61 -55
  147. data/vendor/faiss/faiss/impl/index_read.cpp +547 -479
  148. data/vendor/faiss/faiss/impl/index_write.cpp +497 -407
  149. data/vendor/faiss/faiss/impl/io.cpp +75 -94
  150. data/vendor/faiss/faiss/impl/io.h +31 -41
  151. data/vendor/faiss/faiss/impl/io_macros.h +40 -29
  152. data/vendor/faiss/faiss/impl/lattice_Zn.cpp +137 -186
  153. data/vendor/faiss/faiss/impl/lattice_Zn.h +40 -51
  154. data/vendor/faiss/faiss/impl/platform_macros.h +29 -8
  155. data/vendor/faiss/faiss/impl/pq4_fast_scan.cpp +77 -124
  156. data/vendor/faiss/faiss/impl/pq4_fast_scan.h +39 -48
  157. data/vendor/faiss/faiss/impl/pq4_fast_scan_search_1.cpp +41 -52
  158. data/vendor/faiss/faiss/impl/pq4_fast_scan_search_qbs.cpp +80 -117
  159. data/vendor/faiss/faiss/impl/simd_result_handlers.h +109 -137
  160. data/vendor/faiss/faiss/index_factory.cpp +269 -218
  161. data/vendor/faiss/faiss/index_factory.h +6 -7
  162. data/vendor/faiss/faiss/index_io.h +23 -26
  163. data/vendor/faiss/faiss/invlists/BlockInvertedLists.cpp +67 -75
  164. data/vendor/faiss/faiss/invlists/BlockInvertedLists.h +22 -24
  165. data/vendor/faiss/faiss/invlists/DirectMap.cpp +96 -112
  166. data/vendor/faiss/faiss/invlists/DirectMap.h +29 -33
  167. data/vendor/faiss/faiss/invlists/InvertedLists.cpp +307 -364
  168. data/vendor/faiss/faiss/invlists/InvertedLists.h +151 -151
  169. data/vendor/faiss/faiss/invlists/InvertedListsIOHook.cpp +29 -34
  170. data/vendor/faiss/faiss/invlists/InvertedListsIOHook.h +17 -18
  171. data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.cpp +257 -293
  172. data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.h +50 -45
  173. data/vendor/faiss/faiss/python/python_callbacks.cpp +23 -26
  174. data/vendor/faiss/faiss/python/python_callbacks.h +9 -16
  175. data/vendor/faiss/faiss/utils/AlignedTable.h +79 -44
  176. data/vendor/faiss/faiss/utils/Heap.cpp +40 -48
  177. data/vendor/faiss/faiss/utils/Heap.h +186 -209
  178. data/vendor/faiss/faiss/utils/WorkerThread.cpp +67 -76
  179. data/vendor/faiss/faiss/utils/WorkerThread.h +32 -33
  180. data/vendor/faiss/faiss/utils/distances.cpp +301 -310
  181. data/vendor/faiss/faiss/utils/distances.h +133 -118
  182. data/vendor/faiss/faiss/utils/distances_simd.cpp +456 -516
  183. data/vendor/faiss/faiss/utils/extra_distances-inl.h +117 -0
  184. data/vendor/faiss/faiss/utils/extra_distances.cpp +113 -232
  185. data/vendor/faiss/faiss/utils/extra_distances.h +30 -29
  186. data/vendor/faiss/faiss/utils/hamming-inl.h +260 -209
  187. data/vendor/faiss/faiss/utils/hamming.cpp +375 -469
  188. data/vendor/faiss/faiss/utils/hamming.h +62 -85
  189. data/vendor/faiss/faiss/utils/ordered_key_value.h +16 -18
  190. data/vendor/faiss/faiss/utils/partitioning.cpp +393 -318
  191. data/vendor/faiss/faiss/utils/partitioning.h +26 -21
  192. data/vendor/faiss/faiss/utils/quantize_lut.cpp +78 -66
  193. data/vendor/faiss/faiss/utils/quantize_lut.h +22 -20
  194. data/vendor/faiss/faiss/utils/random.cpp +39 -63
  195. data/vendor/faiss/faiss/utils/random.h +13 -16
  196. data/vendor/faiss/faiss/utils/simdlib.h +4 -2
  197. data/vendor/faiss/faiss/utils/simdlib_avx2.h +88 -85
  198. data/vendor/faiss/faiss/utils/simdlib_emulated.h +226 -165
  199. data/vendor/faiss/faiss/utils/simdlib_neon.h +832 -0
  200. data/vendor/faiss/faiss/utils/utils.cpp +304 -287
  201. data/vendor/faiss/faiss/utils/utils.h +53 -48
  202. metadata +20 -2
@@ -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 { namespace gpu {
21
+ namespace faiss {
22
+ namespace gpu {
23
23
 
24
24
  class FlatIndex;
25
25
 
26
26
  struct GpuIndexFlatConfig : public GpuIndexConfig {
27
- inline GpuIndexFlatConfig()
28
- : useFloat16(false),
29
- storeTransposed(false) {
30
- }
31
-
32
- /// Whether or not data is stored as float16
33
- bool useFloat16;
34
-
35
- /// Whether or not data is stored (transparently) in a transposed
36
- /// layout, enabling use of the NN GEMM call, which is ~10% faster.
37
- /// This will improve the speed of the flat index, but will
38
- /// substantially slow down any add() calls made, as all data must
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
- public:
49
- /// Construct from a pre-existing faiss::IndexFlat instance, copying
50
- /// data over to the given GPU
51
- GpuIndexFlat(GpuResourcesProvider* provider,
52
- const faiss::IndexFlat* index,
53
- GpuIndexFlatConfig config = GpuIndexFlatConfig());
54
-
55
- GpuIndexFlat(std::shared_ptr<GpuResources> resources,
56
- const faiss::IndexFlat* index,
57
- GpuIndexFlatConfig config = GpuIndexFlatConfig());
58
-
59
- /// Construct an empty instance that can be added to
60
- GpuIndexFlat(GpuResourcesProvider* provider,
61
- int dims,
62
- faiss::MetricType metric,
63
- GpuIndexFlatConfig config = GpuIndexFlatConfig());
64
-
65
- GpuIndexFlat(std::shared_ptr<GpuResources> resources,
66
- int dims,
67
- faiss::MetricType metric,
68
- GpuIndexFlatConfig config = GpuIndexFlatConfig());
69
-
70
- ~GpuIndexFlat() override;
71
-
72
- /// Initialize ourselves from the given CPU index; will overwrite
73
- /// all data in ourselves
74
- void copyFrom(const faiss::IndexFlat* index);
75
-
76
- /// Copy ourselves to the given CPU index; will overwrite all data
77
- /// in the index instance
78
- void copyTo(faiss::IndexFlat* index) const;
79
-
80
- /// Returns the number of vectors we contain
81
- size_t getNumVecs() const;
82
-
83
- /// Clears all vectors from this index
84
- void reset() override;
85
-
86
- /// This index is not trained, so this does nothing
87
- void train(Index::idx_t n, const float* x) override;
88
-
89
- /// Overrides to avoid excessive copies
90
- void add(Index::idx_t, const float* x) override;
91
-
92
- /// Reconstruction methods; prefer the batch reconstruct as it will
93
- /// be more efficient
94
- void reconstruct(Index::idx_t key, float* out) const override;
95
-
96
- /// Batch reconstruction method
97
- void reconstruct_n(Index::idx_t i0,
98
- Index::idx_t num,
99
- float* out) const override;
100
-
101
- /// Compute residual
102
- void compute_residual(const float* x,
103
- float* residual,
104
- Index::idx_t key) const override;
105
-
106
- /// Compute residual (batch mode)
107
- void compute_residual_n(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() { return data_.get(); }
114
-
115
- protected:
116
- /// Flat index does not require IDs as there is no storage available for them
117
- bool addImplRequiresIDs_() const override;
118
-
119
- /// Called from GpuIndex for add
120
- void addImpl_(int n,
121
- const float* x,
122
- const Index::idx_t* ids) override;
123
-
124
- /// Called from GpuIndex for search
125
- void searchImpl_(int n,
126
- const float* x,
127
- int k,
128
- float* distances,
129
- Index::idx_t* labels) const override;
130
-
131
- protected:
132
- /// Our configuration options
133
- const GpuIndexFlatConfig flatConfig_;
134
-
135
- /// Holds our GPU data containing the list of vectors
136
- std::unique_ptr<FlatIndex> data_;
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
- public:
144
- /// Construct from a pre-existing faiss::IndexFlatL2 instance, copying
145
- /// data over to the given GPU
146
- GpuIndexFlatL2(GpuResourcesProvider* provider,
147
- faiss::IndexFlatL2* index,
148
- GpuIndexFlatConfig config = GpuIndexFlatConfig());
149
-
150
- GpuIndexFlatL2(std::shared_ptr<GpuResources> resources,
151
- faiss::IndexFlatL2* index,
152
- GpuIndexFlatConfig config = GpuIndexFlatConfig());
153
-
154
- /// Construct an empty instance that can be added to
155
- GpuIndexFlatL2(GpuResourcesProvider* provider,
156
- int dims,
157
- GpuIndexFlatConfig config = GpuIndexFlatConfig());
158
-
159
- GpuIndexFlatL2(std::shared_ptr<GpuResources> resources,
160
- int dims,
161
- GpuIndexFlatConfig config = GpuIndexFlatConfig());
162
-
163
- /// Initialize ourselves from the given CPU index; will overwrite
164
- /// all data in ourselves
165
- void copyFrom(faiss::IndexFlat* index);
166
-
167
- /// Copy ourselves to the given CPU index; will overwrite all data
168
- /// in the index instance
169
- void copyTo(faiss::IndexFlat* index);
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
- public:
177
- /// Construct from a pre-existing faiss::IndexFlatIP instance, copying
178
- /// data over to the given GPU
179
- GpuIndexFlatIP(GpuResourcesProvider* provider,
180
- faiss::IndexFlatIP* index,
181
- GpuIndexFlatConfig config = GpuIndexFlatConfig());
182
-
183
- GpuIndexFlatIP(std::shared_ptr<GpuResources> resources,
184
- faiss::IndexFlatIP* index,
185
- GpuIndexFlatConfig config = GpuIndexFlatConfig());
186
-
187
- /// Construct an empty instance that can be added to
188
- GpuIndexFlatIP(GpuResourcesProvider* provider,
189
- int dims,
190
- GpuIndexFlatConfig config = GpuIndexFlatConfig());
191
-
192
- GpuIndexFlatIP(std::shared_ptr<GpuResources> resources,
193
- int dims,
194
- GpuIndexFlatConfig config = GpuIndexFlatConfig());
195
-
196
- /// Initialize ourselves from the given CPU index; will overwrite
197
- /// all data in ourselves
198
- void copyFrom(faiss::IndexFlat* index);
199
-
200
- /// Copy ourselves to the given CPU index; will overwrite all data
201
- /// in the index instance
202
- void copyTo(faiss::IndexFlat* index);
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
- } } // namespace
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 { struct IndexIVF; }
15
+ namespace faiss {
16
+ struct IndexIVF;
17
+ }
17
18
 
18
- namespace faiss { namespace gpu {
19
+ namespace faiss {
20
+ namespace gpu {
19
21
 
20
22
  class GpuIndexFlat;
21
23
 
22
24
  struct GpuIndexIVFConfig : public GpuIndexConfig {
23
- inline GpuIndexIVFConfig()
24
- : indicesOptions(INDICES_64_BIT) {
25
- }
25
+ inline GpuIndexIVFConfig() : indicesOptions(INDICES_64_BIT) {}
26
26
 
27
- /// Index storage options for the GPU
28
- IndicesOptions indicesOptions;
27
+ /// Index storage options for the GPU
28
+ IndicesOptions indicesOptions;
29
29
 
30
- /// Configuration for the coarse quantizer object
31
- GpuIndexFlatConfig flatConfig;
30
+ /// Configuration for the coarse quantizer object
31
+ GpuIndexFlatConfig flatConfig;
32
32
  };
33
33
 
34
34
  class GpuIndexIVF : public GpuIndex {
35
- public:
36
- GpuIndexIVF(GpuResourcesProvider* provider,
37
- int dims,
38
- faiss::MetricType metric,
39
- float metricArg,
40
- int nlist,
41
- GpuIndexIVFConfig config = GpuIndexIVFConfig());
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
- ~GpuIndexIVF() override;
44
+ ~GpuIndexIVF() override;
44
45
 
45
- private:
46
- /// Shared initialization functions
47
- void init_();
46
+ private:
47
+ /// Shared initialization functions
48
+ void init_();
48
49
 
49
- public:
50
- /// Copy what we need from the CPU equivalent
51
- void copyFrom(const faiss::IndexIVF* index);
50
+ public:
51
+ /// Copy what we need from the CPU equivalent
52
+ void copyFrom(const faiss::IndexIVF* index);
52
53
 
53
- /// Copy what we have to the CPU equivalent
54
- void copyTo(faiss::IndexIVF* index) const;
54
+ /// Copy what we have to the CPU equivalent
55
+ void copyTo(faiss::IndexIVF* index) const;
55
56
 
56
- /// Returns the number of inverted lists we're managing
57
- int getNumLists() const;
57
+ /// Returns the number of inverted lists we're managing
58
+ int getNumLists() const;
58
59
 
59
- /// Returns the number of vectors present in a particular inverted list
60
- virtual int getListLength(int listId) const = 0;
60
+ /// Returns the number of vectors present in a particular inverted list
61
+ virtual int getListLength(int listId) const = 0;
61
62
 
62
- /// Return the encoded vector data contained in a particular inverted list,
63
- /// for debugging purposes.
64
- /// If gpuFormat is true, the data is returned as it is encoded in the
65
- /// GPU-side representation.
66
- /// Otherwise, it is converted to the CPU format.
67
- /// compliant format, while the native GPU format may differ.
68
- virtual std::vector<uint8_t>
69
- getListVectorData(int listId, bool gpuFormat = false) const = 0;
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
- /// Return the vector indices contained in a particular inverted list, for
72
- /// debugging purposes.
73
- virtual std::vector<Index::idx_t> getListIndices(int listId) const = 0;
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
- /// Return the quantizer we're using
76
- GpuIndexFlat* getQuantizer();
77
+ /// Return the quantizer we're using
78
+ GpuIndexFlat* getQuantizer();
77
79
 
78
- /// Sets the number of list probes per query
79
- void setNumProbes(int nprobe);
80
+ /// Sets the number of list probes per query
81
+ void setNumProbes(int nprobe);
80
82
 
81
- /// Returns our current number of list probes per query
82
- int getNumProbes() const;
83
+ /// Returns our current number of list probes per query
84
+ int getNumProbes() const;
83
85
 
84
- protected:
85
- bool addImplRequiresIDs_() const override;
86
- void trainQuantizer_(Index::idx_t n, const float* x);
86
+ protected:
87
+ bool addImplRequiresIDs_() const override;
88
+ void trainQuantizer_(Index::idx_t n, const float* x);
87
89
 
88
- public:
89
- /// Exposing this like the CPU version for manipulation
90
- ClusteringParameters cp;
90
+ public:
91
+ /// Exposing this like the CPU version for manipulation
92
+ ClusteringParameters cp;
91
93
 
92
- /// Exposing this like the CPU version for query
93
- int nlist;
94
+ /// Exposing this like the CPU version for query
95
+ int nlist;
94
96
 
95
- /// Exposing this like the CPU version for manipulation
96
- int nprobe;
97
+ /// Exposing this like the CPU version for manipulation
98
+ int nprobe;
97
99
 
98
- /// Exposeing this like the CPU version for query
99
- GpuIndexFlat* quantizer;
100
+ /// Exposeing this like the CPU version for query
101
+ GpuIndexFlat* quantizer;
100
102
 
101
- protected:
102
- /// Our configuration options
103
- const GpuIndexIVFConfig ivfConfig_;
103
+ protected:
104
+ /// Our configuration options
105
+ const GpuIndexIVFConfig ivfConfig_;
104
106
  };
105
107
 
106
- } } // namespace
108
+ } // namespace gpu
109
+ } // namespace faiss