faiss 0.2.0 → 0.2.1

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 (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