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.
Files changed (215) hide show
  1. checksums.yaml +4 -4
  2. data/CHANGELOG.md +16 -0
  3. data/LICENSE.txt +1 -1
  4. data/README.md +7 -7
  5. data/ext/faiss/extconf.rb +6 -3
  6. data/ext/faiss/numo.hpp +4 -4
  7. data/ext/faiss/utils.cpp +1 -1
  8. data/ext/faiss/utils.h +1 -1
  9. data/lib/faiss/version.rb +1 -1
  10. data/vendor/faiss/faiss/AutoTune.cpp +292 -291
  11. data/vendor/faiss/faiss/AutoTune.h +55 -56
  12. data/vendor/faiss/faiss/Clustering.cpp +365 -194
  13. data/vendor/faiss/faiss/Clustering.h +102 -35
  14. data/vendor/faiss/faiss/IVFlib.cpp +171 -195
  15. data/vendor/faiss/faiss/IVFlib.h +48 -51
  16. data/vendor/faiss/faiss/Index.cpp +85 -103
  17. data/vendor/faiss/faiss/Index.h +54 -48
  18. data/vendor/faiss/faiss/Index2Layer.cpp +126 -224
  19. data/vendor/faiss/faiss/Index2Layer.h +22 -36
  20. data/vendor/faiss/faiss/IndexAdditiveQuantizer.cpp +407 -0
  21. data/vendor/faiss/faiss/IndexAdditiveQuantizer.h +195 -0
  22. data/vendor/faiss/faiss/IndexBinary.cpp +45 -37
  23. data/vendor/faiss/faiss/IndexBinary.h +140 -132
  24. data/vendor/faiss/faiss/IndexBinaryFlat.cpp +73 -53
  25. data/vendor/faiss/faiss/IndexBinaryFlat.h +29 -24
  26. data/vendor/faiss/faiss/IndexBinaryFromFloat.cpp +46 -43
  27. data/vendor/faiss/faiss/IndexBinaryFromFloat.h +16 -15
  28. data/vendor/faiss/faiss/IndexBinaryHNSW.cpp +215 -232
  29. data/vendor/faiss/faiss/IndexBinaryHNSW.h +25 -24
  30. data/vendor/faiss/faiss/IndexBinaryHash.cpp +182 -177
  31. data/vendor/faiss/faiss/IndexBinaryHash.h +41 -34
  32. data/vendor/faiss/faiss/IndexBinaryIVF.cpp +489 -461
  33. data/vendor/faiss/faiss/IndexBinaryIVF.h +97 -68
  34. data/vendor/faiss/faiss/IndexFlat.cpp +115 -176
  35. data/vendor/faiss/faiss/IndexFlat.h +42 -59
  36. data/vendor/faiss/faiss/IndexFlatCodes.cpp +67 -0
  37. data/vendor/faiss/faiss/IndexFlatCodes.h +47 -0
  38. data/vendor/faiss/faiss/IndexHNSW.cpp +372 -348
  39. data/vendor/faiss/faiss/IndexHNSW.h +57 -41
  40. data/vendor/faiss/faiss/IndexIVF.cpp +545 -453
  41. data/vendor/faiss/faiss/IndexIVF.h +169 -118
  42. data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.cpp +316 -0
  43. data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.h +121 -0
  44. data/vendor/faiss/faiss/IndexIVFFlat.cpp +247 -252
  45. data/vendor/faiss/faiss/IndexIVFFlat.h +48 -51
  46. data/vendor/faiss/faiss/IndexIVFPQ.cpp +459 -517
  47. data/vendor/faiss/faiss/IndexIVFPQ.h +75 -67
  48. data/vendor/faiss/faiss/IndexIVFPQFastScan.cpp +406 -372
  49. data/vendor/faiss/faiss/IndexIVFPQFastScan.h +82 -57
  50. data/vendor/faiss/faiss/IndexIVFPQR.cpp +104 -102
  51. data/vendor/faiss/faiss/IndexIVFPQR.h +33 -28
  52. data/vendor/faiss/faiss/IndexIVFSpectralHash.cpp +163 -150
  53. data/vendor/faiss/faiss/IndexIVFSpectralHash.h +38 -25
  54. data/vendor/faiss/faiss/IndexLSH.cpp +66 -113
  55. data/vendor/faiss/faiss/IndexLSH.h +20 -38
  56. data/vendor/faiss/faiss/IndexLattice.cpp +42 -56
  57. data/vendor/faiss/faiss/IndexLattice.h +11 -16
  58. data/vendor/faiss/faiss/IndexNNDescent.cpp +229 -0
  59. data/vendor/faiss/faiss/IndexNNDescent.h +72 -0
  60. data/vendor/faiss/faiss/IndexNSG.cpp +301 -0
  61. data/vendor/faiss/faiss/IndexNSG.h +85 -0
  62. data/vendor/faiss/faiss/IndexPQ.cpp +387 -495
  63. data/vendor/faiss/faiss/IndexPQ.h +64 -82
  64. data/vendor/faiss/faiss/IndexPQFastScan.cpp +143 -170
  65. data/vendor/faiss/faiss/IndexPQFastScan.h +46 -32
  66. data/vendor/faiss/faiss/IndexPreTransform.cpp +120 -150
  67. data/vendor/faiss/faiss/IndexPreTransform.h +33 -36
  68. data/vendor/faiss/faiss/IndexRefine.cpp +139 -127
  69. data/vendor/faiss/faiss/IndexRefine.h +32 -23
  70. data/vendor/faiss/faiss/IndexReplicas.cpp +147 -153
  71. data/vendor/faiss/faiss/IndexReplicas.h +62 -56
  72. data/vendor/faiss/faiss/IndexScalarQuantizer.cpp +111 -172
  73. data/vendor/faiss/faiss/IndexScalarQuantizer.h +41 -59
  74. data/vendor/faiss/faiss/IndexShards.cpp +256 -240
  75. data/vendor/faiss/faiss/IndexShards.h +85 -73
  76. data/vendor/faiss/faiss/MatrixStats.cpp +112 -97
  77. data/vendor/faiss/faiss/MatrixStats.h +7 -10
  78. data/vendor/faiss/faiss/MetaIndexes.cpp +135 -157
  79. data/vendor/faiss/faiss/MetaIndexes.h +40 -34
  80. data/vendor/faiss/faiss/MetricType.h +7 -7
  81. data/vendor/faiss/faiss/VectorTransform.cpp +654 -475
  82. data/vendor/faiss/faiss/VectorTransform.h +64 -89
  83. data/vendor/faiss/faiss/clone_index.cpp +78 -73
  84. data/vendor/faiss/faiss/clone_index.h +4 -9
  85. data/vendor/faiss/faiss/gpu/GpuAutoTune.cpp +33 -38
  86. data/vendor/faiss/faiss/gpu/GpuAutoTune.h +11 -9
  87. data/vendor/faiss/faiss/gpu/GpuCloner.cpp +198 -171
  88. data/vendor/faiss/faiss/gpu/GpuCloner.h +53 -35
  89. data/vendor/faiss/faiss/gpu/GpuClonerOptions.cpp +12 -14
  90. data/vendor/faiss/faiss/gpu/GpuClonerOptions.h +27 -25
  91. data/vendor/faiss/faiss/gpu/GpuDistance.h +116 -112
  92. data/vendor/faiss/faiss/gpu/GpuFaissAssert.h +1 -2
  93. data/vendor/faiss/faiss/gpu/GpuIcmEncoder.h +60 -0
  94. data/vendor/faiss/faiss/gpu/GpuIndex.h +134 -137
  95. data/vendor/faiss/faiss/gpu/GpuIndexBinaryFlat.h +76 -73
  96. data/vendor/faiss/faiss/gpu/GpuIndexFlat.h +173 -162
  97. data/vendor/faiss/faiss/gpu/GpuIndexIVF.h +67 -64
  98. data/vendor/faiss/faiss/gpu/GpuIndexIVFFlat.h +89 -86
  99. data/vendor/faiss/faiss/gpu/GpuIndexIVFPQ.h +150 -141
  100. data/vendor/faiss/faiss/gpu/GpuIndexIVFScalarQuantizer.h +101 -103
  101. data/vendor/faiss/faiss/gpu/GpuIndicesOptions.h +17 -16
  102. data/vendor/faiss/faiss/gpu/GpuResources.cpp +116 -128
  103. data/vendor/faiss/faiss/gpu/GpuResources.h +182 -186
  104. data/vendor/faiss/faiss/gpu/StandardGpuResources.cpp +433 -422
  105. data/vendor/faiss/faiss/gpu/StandardGpuResources.h +131 -130
  106. data/vendor/faiss/faiss/gpu/impl/InterleavedCodes.cpp +468 -456
  107. data/vendor/faiss/faiss/gpu/impl/InterleavedCodes.h +25 -19
  108. data/vendor/faiss/faiss/gpu/impl/RemapIndices.cpp +22 -20
  109. data/vendor/faiss/faiss/gpu/impl/RemapIndices.h +9 -8
  110. data/vendor/faiss/faiss/gpu/perf/IndexWrapper-inl.h +39 -44
  111. data/vendor/faiss/faiss/gpu/perf/IndexWrapper.h +16 -14
  112. data/vendor/faiss/faiss/gpu/perf/PerfClustering.cpp +77 -71
  113. data/vendor/faiss/faiss/gpu/perf/PerfIVFPQAdd.cpp +109 -88
  114. data/vendor/faiss/faiss/gpu/perf/WriteIndex.cpp +75 -64
  115. data/vendor/faiss/faiss/gpu/test/TestCodePacking.cpp +230 -215
  116. data/vendor/faiss/faiss/gpu/test/TestGpuIndexBinaryFlat.cpp +80 -86
  117. data/vendor/faiss/faiss/gpu/test/TestGpuIndexFlat.cpp +284 -277
  118. data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFFlat.cpp +416 -416
  119. data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFPQ.cpp +611 -517
  120. data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFScalarQuantizer.cpp +166 -164
  121. data/vendor/faiss/faiss/gpu/test/TestGpuMemoryException.cpp +61 -53
  122. data/vendor/faiss/faiss/gpu/test/TestUtils.cpp +274 -238
  123. data/vendor/faiss/faiss/gpu/test/TestUtils.h +73 -57
  124. data/vendor/faiss/faiss/gpu/test/demo_ivfpq_indexing_gpu.cpp +47 -50
  125. data/vendor/faiss/faiss/gpu/utils/DeviceUtils.h +79 -72
  126. data/vendor/faiss/faiss/gpu/utils/StackDeviceMemory.cpp +140 -146
  127. data/vendor/faiss/faiss/gpu/utils/StackDeviceMemory.h +69 -71
  128. data/vendor/faiss/faiss/gpu/utils/StaticUtils.h +21 -16
  129. data/vendor/faiss/faiss/gpu/utils/Timer.cpp +25 -29
  130. data/vendor/faiss/faiss/gpu/utils/Timer.h +30 -29
  131. data/vendor/faiss/faiss/impl/AdditiveQuantizer.cpp +503 -0
  132. data/vendor/faiss/faiss/impl/AdditiveQuantizer.h +175 -0
  133. data/vendor/faiss/faiss/impl/AuxIndexStructures.cpp +90 -120
  134. data/vendor/faiss/faiss/impl/AuxIndexStructures.h +81 -65
  135. data/vendor/faiss/faiss/impl/FaissAssert.h +73 -58
  136. data/vendor/faiss/faiss/impl/FaissException.cpp +56 -48
  137. data/vendor/faiss/faiss/impl/FaissException.h +41 -29
  138. data/vendor/faiss/faiss/impl/HNSW.cpp +606 -617
  139. data/vendor/faiss/faiss/impl/HNSW.h +179 -200
  140. data/vendor/faiss/faiss/impl/LocalSearchQuantizer.cpp +855 -0
  141. data/vendor/faiss/faiss/impl/LocalSearchQuantizer.h +244 -0
  142. data/vendor/faiss/faiss/impl/NNDescent.cpp +487 -0
  143. data/vendor/faiss/faiss/impl/NNDescent.h +154 -0
  144. data/vendor/faiss/faiss/impl/NSG.cpp +679 -0
  145. data/vendor/faiss/faiss/impl/NSG.h +199 -0
  146. data/vendor/faiss/faiss/impl/PolysemousTraining.cpp +484 -454
  147. data/vendor/faiss/faiss/impl/PolysemousTraining.h +52 -55
  148. data/vendor/faiss/faiss/impl/ProductQuantizer-inl.h +26 -47
  149. data/vendor/faiss/faiss/impl/ProductQuantizer.cpp +469 -459
  150. data/vendor/faiss/faiss/impl/ProductQuantizer.h +76 -87
  151. data/vendor/faiss/faiss/impl/ResidualQuantizer.cpp +758 -0
  152. data/vendor/faiss/faiss/impl/ResidualQuantizer.h +188 -0
  153. data/vendor/faiss/faiss/impl/ResultHandler.h +96 -132
  154. data/vendor/faiss/faiss/impl/ScalarQuantizer.cpp +647 -707
  155. data/vendor/faiss/faiss/impl/ScalarQuantizer.h +48 -46
  156. data/vendor/faiss/faiss/impl/ThreadedIndex-inl.h +129 -131
  157. data/vendor/faiss/faiss/impl/ThreadedIndex.h +61 -55
  158. data/vendor/faiss/faiss/impl/index_read.cpp +631 -480
  159. data/vendor/faiss/faiss/impl/index_write.cpp +547 -407
  160. data/vendor/faiss/faiss/impl/io.cpp +76 -95
  161. data/vendor/faiss/faiss/impl/io.h +31 -41
  162. data/vendor/faiss/faiss/impl/io_macros.h +60 -29
  163. data/vendor/faiss/faiss/impl/kmeans1d.cpp +301 -0
  164. data/vendor/faiss/faiss/impl/kmeans1d.h +48 -0
  165. data/vendor/faiss/faiss/impl/lattice_Zn.cpp +137 -186
  166. data/vendor/faiss/faiss/impl/lattice_Zn.h +40 -51
  167. data/vendor/faiss/faiss/impl/platform_macros.h +29 -8
  168. data/vendor/faiss/faiss/impl/pq4_fast_scan.cpp +77 -124
  169. data/vendor/faiss/faiss/impl/pq4_fast_scan.h +39 -48
  170. data/vendor/faiss/faiss/impl/pq4_fast_scan_search_1.cpp +41 -52
  171. data/vendor/faiss/faiss/impl/pq4_fast_scan_search_qbs.cpp +80 -117
  172. data/vendor/faiss/faiss/impl/simd_result_handlers.h +109 -137
  173. data/vendor/faiss/faiss/index_factory.cpp +619 -397
  174. data/vendor/faiss/faiss/index_factory.h +8 -6
  175. data/vendor/faiss/faiss/index_io.h +23 -26
  176. data/vendor/faiss/faiss/invlists/BlockInvertedLists.cpp +67 -75
  177. data/vendor/faiss/faiss/invlists/BlockInvertedLists.h +22 -24
  178. data/vendor/faiss/faiss/invlists/DirectMap.cpp +96 -112
  179. data/vendor/faiss/faiss/invlists/DirectMap.h +29 -33
  180. data/vendor/faiss/faiss/invlists/InvertedLists.cpp +307 -364
  181. data/vendor/faiss/faiss/invlists/InvertedLists.h +151 -151
  182. data/vendor/faiss/faiss/invlists/InvertedListsIOHook.cpp +29 -34
  183. data/vendor/faiss/faiss/invlists/InvertedListsIOHook.h +17 -18
  184. data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.cpp +257 -293
  185. data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.h +50 -45
  186. data/vendor/faiss/faiss/python/python_callbacks.cpp +23 -26
  187. data/vendor/faiss/faiss/python/python_callbacks.h +9 -16
  188. data/vendor/faiss/faiss/utils/AlignedTable.h +79 -44
  189. data/vendor/faiss/faiss/utils/Heap.cpp +40 -48
  190. data/vendor/faiss/faiss/utils/Heap.h +186 -209
  191. data/vendor/faiss/faiss/utils/WorkerThread.cpp +67 -76
  192. data/vendor/faiss/faiss/utils/WorkerThread.h +32 -33
  193. data/vendor/faiss/faiss/utils/distances.cpp +305 -312
  194. data/vendor/faiss/faiss/utils/distances.h +170 -122
  195. data/vendor/faiss/faiss/utils/distances_simd.cpp +498 -508
  196. data/vendor/faiss/faiss/utils/extra_distances-inl.h +117 -0
  197. data/vendor/faiss/faiss/utils/extra_distances.cpp +113 -232
  198. data/vendor/faiss/faiss/utils/extra_distances.h +30 -29
  199. data/vendor/faiss/faiss/utils/hamming-inl.h +260 -209
  200. data/vendor/faiss/faiss/utils/hamming.cpp +375 -469
  201. data/vendor/faiss/faiss/utils/hamming.h +62 -85
  202. data/vendor/faiss/faiss/utils/ordered_key_value.h +16 -18
  203. data/vendor/faiss/faiss/utils/partitioning.cpp +393 -318
  204. data/vendor/faiss/faiss/utils/partitioning.h +26 -21
  205. data/vendor/faiss/faiss/utils/quantize_lut.cpp +78 -66
  206. data/vendor/faiss/faiss/utils/quantize_lut.h +22 -20
  207. data/vendor/faiss/faiss/utils/random.cpp +39 -63
  208. data/vendor/faiss/faiss/utils/random.h +13 -16
  209. data/vendor/faiss/faiss/utils/simdlib.h +4 -2
  210. data/vendor/faiss/faiss/utils/simdlib_avx2.h +88 -85
  211. data/vendor/faiss/faiss/utils/simdlib_emulated.h +226 -165
  212. data/vendor/faiss/faiss/utils/simdlib_neon.h +832 -0
  213. data/vendor/faiss/faiss/utils/utils.cpp +304 -287
  214. data/vendor/faiss/faiss/utils/utils.h +54 -49
  215. 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 { 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