faiss 0.1.5 → 0.2.2

Sign up to get free protection for your applications and to get access to all the features.
Files changed (219) hide show
  1. checksums.yaml +4 -4
  2. data/CHANGELOG.md +24 -0
  3. data/README.md +12 -0
  4. data/ext/faiss/ext.cpp +1 -1
  5. data/ext/faiss/extconf.rb +6 -2
  6. data/ext/faiss/index.cpp +114 -43
  7. data/ext/faiss/index_binary.cpp +24 -30
  8. data/ext/faiss/kmeans.cpp +20 -16
  9. data/ext/faiss/numo.hpp +867 -0
  10. data/ext/faiss/pca_matrix.cpp +13 -14
  11. data/ext/faiss/product_quantizer.cpp +23 -24
  12. data/ext/faiss/utils.cpp +10 -37
  13. data/ext/faiss/utils.h +2 -13
  14. data/lib/faiss.rb +0 -5
  15. data/lib/faiss/version.rb +1 -1
  16. data/vendor/faiss/faiss/AutoTune.cpp +292 -291
  17. data/vendor/faiss/faiss/AutoTune.h +55 -56
  18. data/vendor/faiss/faiss/Clustering.cpp +334 -195
  19. data/vendor/faiss/faiss/Clustering.h +88 -35
  20. data/vendor/faiss/faiss/IVFlib.cpp +171 -195
  21. data/vendor/faiss/faiss/IVFlib.h +48 -51
  22. data/vendor/faiss/faiss/Index.cpp +85 -103
  23. data/vendor/faiss/faiss/Index.h +54 -48
  24. data/vendor/faiss/faiss/Index2Layer.cpp +139 -164
  25. data/vendor/faiss/faiss/Index2Layer.h +22 -22
  26. data/vendor/faiss/faiss/IndexBinary.cpp +45 -37
  27. data/vendor/faiss/faiss/IndexBinary.h +140 -132
  28. data/vendor/faiss/faiss/IndexBinaryFlat.cpp +73 -53
  29. data/vendor/faiss/faiss/IndexBinaryFlat.h +29 -24
  30. data/vendor/faiss/faiss/IndexBinaryFromFloat.cpp +46 -43
  31. data/vendor/faiss/faiss/IndexBinaryFromFloat.h +16 -15
  32. data/vendor/faiss/faiss/IndexBinaryHNSW.cpp +215 -232
  33. data/vendor/faiss/faiss/IndexBinaryHNSW.h +25 -24
  34. data/vendor/faiss/faiss/IndexBinaryHash.cpp +182 -177
  35. data/vendor/faiss/faiss/IndexBinaryHash.h +41 -34
  36. data/vendor/faiss/faiss/IndexBinaryIVF.cpp +489 -461
  37. data/vendor/faiss/faiss/IndexBinaryIVF.h +97 -68
  38. data/vendor/faiss/faiss/IndexFlat.cpp +116 -147
  39. data/vendor/faiss/faiss/IndexFlat.h +35 -46
  40. data/vendor/faiss/faiss/IndexHNSW.cpp +372 -348
  41. data/vendor/faiss/faiss/IndexHNSW.h +57 -41
  42. data/vendor/faiss/faiss/IndexIVF.cpp +474 -454
  43. data/vendor/faiss/faiss/IndexIVF.h +146 -113
  44. data/vendor/faiss/faiss/IndexIVFFlat.cpp +248 -250
  45. data/vendor/faiss/faiss/IndexIVFFlat.h +48 -51
  46. data/vendor/faiss/faiss/IndexIVFPQ.cpp +457 -516
  47. data/vendor/faiss/faiss/IndexIVFPQ.h +74 -66
  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 +125 -133
  53. data/vendor/faiss/faiss/IndexIVFSpectralHash.h +19 -21
  54. data/vendor/faiss/faiss/IndexLSH.cpp +75 -96
  55. data/vendor/faiss/faiss/IndexLSH.h +21 -26
  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 +231 -0
  59. data/vendor/faiss/faiss/IndexNNDescent.h +72 -0
  60. data/vendor/faiss/faiss/IndexNSG.cpp +303 -0
  61. data/vendor/faiss/faiss/IndexNSG.h +85 -0
  62. data/vendor/faiss/faiss/IndexPQ.cpp +405 -464
  63. data/vendor/faiss/faiss/IndexPQ.h +64 -67
  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 +115 -131
  69. data/vendor/faiss/faiss/IndexRefine.h +22 -23
  70. data/vendor/faiss/faiss/IndexReplicas.cpp +147 -153
  71. data/vendor/faiss/faiss/IndexReplicas.h +62 -56
  72. data/vendor/faiss/faiss/IndexResidual.cpp +291 -0
  73. data/vendor/faiss/faiss/IndexResidual.h +152 -0
  74. data/vendor/faiss/faiss/IndexScalarQuantizer.cpp +120 -155
  75. data/vendor/faiss/faiss/IndexScalarQuantizer.h +41 -45
  76. data/vendor/faiss/faiss/IndexShards.cpp +256 -240
  77. data/vendor/faiss/faiss/IndexShards.h +85 -73
  78. data/vendor/faiss/faiss/MatrixStats.cpp +112 -97
  79. data/vendor/faiss/faiss/MatrixStats.h +7 -10
  80. data/vendor/faiss/faiss/MetaIndexes.cpp +135 -157
  81. data/vendor/faiss/faiss/MetaIndexes.h +40 -34
  82. data/vendor/faiss/faiss/MetricType.h +7 -7
  83. data/vendor/faiss/faiss/VectorTransform.cpp +652 -474
  84. data/vendor/faiss/faiss/VectorTransform.h +61 -89
  85. data/vendor/faiss/faiss/clone_index.cpp +77 -73
  86. data/vendor/faiss/faiss/clone_index.h +4 -9
  87. data/vendor/faiss/faiss/gpu/GpuAutoTune.cpp +33 -38
  88. data/vendor/faiss/faiss/gpu/GpuAutoTune.h +11 -9
  89. data/vendor/faiss/faiss/gpu/GpuCloner.cpp +197 -170
  90. data/vendor/faiss/faiss/gpu/GpuCloner.h +53 -35
  91. data/vendor/faiss/faiss/gpu/GpuClonerOptions.cpp +12 -14
  92. data/vendor/faiss/faiss/gpu/GpuClonerOptions.h +27 -25
  93. data/vendor/faiss/faiss/gpu/GpuDistance.h +116 -112
  94. data/vendor/faiss/faiss/gpu/GpuFaissAssert.h +1 -2
  95. data/vendor/faiss/faiss/gpu/GpuIndex.h +134 -137
  96. data/vendor/faiss/faiss/gpu/GpuIndexBinaryFlat.h +76 -73
  97. data/vendor/faiss/faiss/gpu/GpuIndexFlat.h +173 -162
  98. data/vendor/faiss/faiss/gpu/GpuIndexIVF.h +67 -64
  99. data/vendor/faiss/faiss/gpu/GpuIndexIVFFlat.h +89 -86
  100. data/vendor/faiss/faiss/gpu/GpuIndexIVFPQ.h +150 -141
  101. data/vendor/faiss/faiss/gpu/GpuIndexIVFScalarQuantizer.h +101 -103
  102. data/vendor/faiss/faiss/gpu/GpuIndicesOptions.h +17 -16
  103. data/vendor/faiss/faiss/gpu/GpuResources.cpp +116 -128
  104. data/vendor/faiss/faiss/gpu/GpuResources.h +182 -186
  105. data/vendor/faiss/faiss/gpu/StandardGpuResources.cpp +433 -422
  106. data/vendor/faiss/faiss/gpu/StandardGpuResources.h +131 -130
  107. data/vendor/faiss/faiss/gpu/impl/InterleavedCodes.cpp +468 -456
  108. data/vendor/faiss/faiss/gpu/impl/InterleavedCodes.h +25 -19
  109. data/vendor/faiss/faiss/gpu/impl/RemapIndices.cpp +22 -20
  110. data/vendor/faiss/faiss/gpu/impl/RemapIndices.h +9 -8
  111. data/vendor/faiss/faiss/gpu/perf/IndexWrapper-inl.h +39 -44
  112. data/vendor/faiss/faiss/gpu/perf/IndexWrapper.h +16 -14
  113. data/vendor/faiss/faiss/gpu/perf/PerfClustering.cpp +77 -71
  114. data/vendor/faiss/faiss/gpu/perf/PerfIVFPQAdd.cpp +109 -88
  115. data/vendor/faiss/faiss/gpu/perf/WriteIndex.cpp +75 -64
  116. data/vendor/faiss/faiss/gpu/test/TestCodePacking.cpp +230 -215
  117. data/vendor/faiss/faiss/gpu/test/TestGpuIndexBinaryFlat.cpp +80 -86
  118. data/vendor/faiss/faiss/gpu/test/TestGpuIndexFlat.cpp +284 -277
  119. data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFFlat.cpp +416 -416
  120. data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFPQ.cpp +611 -517
  121. data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFScalarQuantizer.cpp +166 -164
  122. data/vendor/faiss/faiss/gpu/test/TestGpuMemoryException.cpp +61 -53
  123. data/vendor/faiss/faiss/gpu/test/TestUtils.cpp +274 -238
  124. data/vendor/faiss/faiss/gpu/test/TestUtils.h +73 -57
  125. data/vendor/faiss/faiss/gpu/test/demo_ivfpq_indexing_gpu.cpp +47 -50
  126. data/vendor/faiss/faiss/gpu/utils/DeviceUtils.h +79 -72
  127. data/vendor/faiss/faiss/gpu/utils/StackDeviceMemory.cpp +140 -146
  128. data/vendor/faiss/faiss/gpu/utils/StackDeviceMemory.h +69 -71
  129. data/vendor/faiss/faiss/gpu/utils/StaticUtils.h +21 -16
  130. data/vendor/faiss/faiss/gpu/utils/Timer.cpp +25 -29
  131. data/vendor/faiss/faiss/gpu/utils/Timer.h +30 -29
  132. data/vendor/faiss/faiss/impl/AdditiveQuantizer.cpp +270 -0
  133. data/vendor/faiss/faiss/impl/AdditiveQuantizer.h +115 -0
  134. data/vendor/faiss/faiss/impl/AuxIndexStructures.cpp +90 -120
  135. data/vendor/faiss/faiss/impl/AuxIndexStructures.h +81 -65
  136. data/vendor/faiss/faiss/impl/FaissAssert.h +73 -58
  137. data/vendor/faiss/faiss/impl/FaissException.cpp +56 -48
  138. data/vendor/faiss/faiss/impl/FaissException.h +41 -29
  139. data/vendor/faiss/faiss/impl/HNSW.cpp +595 -611
  140. data/vendor/faiss/faiss/impl/HNSW.h +179 -200
  141. data/vendor/faiss/faiss/impl/LocalSearchQuantizer.cpp +672 -0
  142. data/vendor/faiss/faiss/impl/LocalSearchQuantizer.h +172 -0
  143. data/vendor/faiss/faiss/impl/NNDescent.cpp +487 -0
  144. data/vendor/faiss/faiss/impl/NNDescent.h +154 -0
  145. data/vendor/faiss/faiss/impl/NSG.cpp +682 -0
  146. data/vendor/faiss/faiss/impl/NSG.h +199 -0
  147. data/vendor/faiss/faiss/impl/PolysemousTraining.cpp +484 -454
  148. data/vendor/faiss/faiss/impl/PolysemousTraining.h +52 -55
  149. data/vendor/faiss/faiss/impl/ProductQuantizer-inl.h +26 -47
  150. data/vendor/faiss/faiss/impl/ProductQuantizer.cpp +469 -459
  151. data/vendor/faiss/faiss/impl/ProductQuantizer.h +76 -87
  152. data/vendor/faiss/faiss/impl/ResidualQuantizer.cpp +448 -0
  153. data/vendor/faiss/faiss/impl/ResidualQuantizer.h +130 -0
  154. data/vendor/faiss/faiss/impl/ResultHandler.h +96 -132
  155. data/vendor/faiss/faiss/impl/ScalarQuantizer.cpp +648 -701
  156. data/vendor/faiss/faiss/impl/ScalarQuantizer.h +48 -46
  157. data/vendor/faiss/faiss/impl/ThreadedIndex-inl.h +129 -131
  158. data/vendor/faiss/faiss/impl/ThreadedIndex.h +61 -55
  159. data/vendor/faiss/faiss/impl/index_read.cpp +547 -479
  160. data/vendor/faiss/faiss/impl/index_write.cpp +497 -407
  161. data/vendor/faiss/faiss/impl/io.cpp +75 -94
  162. data/vendor/faiss/faiss/impl/io.h +31 -41
  163. data/vendor/faiss/faiss/impl/io_macros.h +40 -29
  164. data/vendor/faiss/faiss/impl/lattice_Zn.cpp +137 -186
  165. data/vendor/faiss/faiss/impl/lattice_Zn.h +40 -51
  166. data/vendor/faiss/faiss/impl/platform_macros.h +29 -8
  167. data/vendor/faiss/faiss/impl/pq4_fast_scan.cpp +77 -124
  168. data/vendor/faiss/faiss/impl/pq4_fast_scan.h +39 -48
  169. data/vendor/faiss/faiss/impl/pq4_fast_scan_search_1.cpp +41 -52
  170. data/vendor/faiss/faiss/impl/pq4_fast_scan_search_qbs.cpp +80 -117
  171. data/vendor/faiss/faiss/impl/simd_result_handlers.h +109 -137
  172. data/vendor/faiss/faiss/index_factory.cpp +269 -218
  173. data/vendor/faiss/faiss/index_factory.h +6 -7
  174. data/vendor/faiss/faiss/index_io.h +23 -26
  175. data/vendor/faiss/faiss/invlists/BlockInvertedLists.cpp +67 -75
  176. data/vendor/faiss/faiss/invlists/BlockInvertedLists.h +22 -24
  177. data/vendor/faiss/faiss/invlists/DirectMap.cpp +96 -112
  178. data/vendor/faiss/faiss/invlists/DirectMap.h +29 -33
  179. data/vendor/faiss/faiss/invlists/InvertedLists.cpp +307 -364
  180. data/vendor/faiss/faiss/invlists/InvertedLists.h +151 -151
  181. data/vendor/faiss/faiss/invlists/InvertedListsIOHook.cpp +29 -34
  182. data/vendor/faiss/faiss/invlists/InvertedListsIOHook.h +17 -18
  183. data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.cpp +257 -293
  184. data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.h +50 -45
  185. data/vendor/faiss/faiss/python/python_callbacks.cpp +23 -26
  186. data/vendor/faiss/faiss/python/python_callbacks.h +9 -16
  187. data/vendor/faiss/faiss/utils/AlignedTable.h +79 -44
  188. data/vendor/faiss/faiss/utils/Heap.cpp +40 -48
  189. data/vendor/faiss/faiss/utils/Heap.h +186 -209
  190. data/vendor/faiss/faiss/utils/WorkerThread.cpp +67 -76
  191. data/vendor/faiss/faiss/utils/WorkerThread.h +32 -33
  192. data/vendor/faiss/faiss/utils/distances.cpp +301 -310
  193. data/vendor/faiss/faiss/utils/distances.h +133 -118
  194. data/vendor/faiss/faiss/utils/distances_simd.cpp +456 -516
  195. data/vendor/faiss/faiss/utils/extra_distances-inl.h +117 -0
  196. data/vendor/faiss/faiss/utils/extra_distances.cpp +113 -232
  197. data/vendor/faiss/faiss/utils/extra_distances.h +30 -29
  198. data/vendor/faiss/faiss/utils/hamming-inl.h +260 -209
  199. data/vendor/faiss/faiss/utils/hamming.cpp +375 -469
  200. data/vendor/faiss/faiss/utils/hamming.h +62 -85
  201. data/vendor/faiss/faiss/utils/ordered_key_value.h +16 -18
  202. data/vendor/faiss/faiss/utils/partitioning.cpp +393 -318
  203. data/vendor/faiss/faiss/utils/partitioning.h +26 -21
  204. data/vendor/faiss/faiss/utils/quantize_lut.cpp +78 -66
  205. data/vendor/faiss/faiss/utils/quantize_lut.h +22 -20
  206. data/vendor/faiss/faiss/utils/random.cpp +39 -63
  207. data/vendor/faiss/faiss/utils/random.h +13 -16
  208. data/vendor/faiss/faiss/utils/simdlib.h +4 -2
  209. data/vendor/faiss/faiss/utils/simdlib_avx2.h +88 -85
  210. data/vendor/faiss/faiss/utils/simdlib_emulated.h +226 -165
  211. data/vendor/faiss/faiss/utils/simdlib_neon.h +832 -0
  212. data/vendor/faiss/faiss/utils/utils.cpp +304 -287
  213. data/vendor/faiss/faiss/utils/utils.h +53 -48
  214. metadata +24 -10
  215. data/lib/faiss/index.rb +0 -20
  216. data/lib/faiss/index_binary.rb +0 -20
  217. data/lib/faiss/kmeans.rb +0 -15
  218. data/lib/faiss/pca_matrix.rb +0 -15
  219. data/lib/faiss/product_quantizer.rb +0 -22
@@ -5,12 +5,11 @@
5
5
  * LICENSE file in the root directory of this source tree.
6
6
  */
7
7
 
8
-
9
8
  #ifndef GPU_FAISS_ASSERT_INCLUDED
10
9
  #define GPU_FAISS_ASSERT_INCLUDED
11
10
 
12
- #include <faiss/impl/FaissAssert.h>
13
11
  #include <cuda.h>
12
+ #include <faiss/impl/FaissAssert.h>
14
13
 
15
14
  ///
16
15
  /// Assertions
@@ -5,152 +5,149 @@
5
5
  * LICENSE file in the root directory of this source tree.
6
6
  */
7
7
 
8
-
9
8
  #pragma once
10
9
 
11
10
  #include <faiss/Index.h>
12
11
  #include <faiss/gpu/GpuResources.h>
13
12
 
14
- namespace faiss { namespace gpu {
13
+ namespace faiss {
14
+ namespace gpu {
15
15
 
16
16
  struct GpuIndexConfig {
17
- inline GpuIndexConfig()
18
- : device(0),
19
- memorySpace(MemorySpace::Device) {
20
- }
21
-
22
- /// GPU device on which the index is resident
23
- int device;
24
-
25
- /// What memory space to use for primary storage.
26
- /// On Pascal and above (CC 6+) architectures, allows GPUs to use
27
- /// more memory than is available on the GPU.
28
- MemorySpace memorySpace;
17
+ inline GpuIndexConfig() : device(0), memorySpace(MemorySpace::Device) {}
18
+
19
+ /// GPU device on which the index is resident
20
+ int device;
21
+
22
+ /// What memory space to use for primary storage.
23
+ /// On Pascal and above (CC 6+) architectures, allows GPUs to use
24
+ /// more memory than is available on the GPU.
25
+ MemorySpace memorySpace;
29
26
  };
30
27
 
31
28
  class GpuIndex : public faiss::Index {
32
- public:
33
- GpuIndex(std::shared_ptr<GpuResources> resources,
34
- int dims,
35
- faiss::MetricType metric,
36
- float metricArg,
37
- GpuIndexConfig config);
38
-
39
- /// Returns the device that this index is resident on
40
- int getDevice() const;
41
-
42
- /// Returns a reference to our GpuResources object that manages memory, stream
43
- /// and handle resources on the GPU
44
- std::shared_ptr<GpuResources> getResources();
45
-
46
- /// Set the minimum data size for searches (in MiB) for which we use
47
- /// CPU -> GPU paging
48
- void setMinPagingSize(size_t size);
49
-
50
- /// Returns the current minimum data size for paged searches
51
- size_t getMinPagingSize() const;
52
-
53
- /// `x` can be resident on the CPU or any GPU; copies are performed
54
- /// as needed
55
- /// Handles paged adds if the add set is too large; calls addInternal_
56
- void add(Index::idx_t, const float* x) override;
57
-
58
- /// `x` and `ids` can be resident on the CPU or any GPU; copies are
59
- /// performed as needed
60
- /// Handles paged adds if the add set is too large; calls addInternal_
61
- void add_with_ids(Index::idx_t n,
62
- const float* x,
63
- const Index::idx_t* ids) override;
64
-
65
- /// `x` and `labels` can be resident on the CPU or any GPU; copies are
66
- /// performed as needed
67
- void assign(Index::idx_t n,
68
- const float* x,
69
- Index::idx_t* labels,
70
- Index::idx_t k = 1) const override;
71
-
72
- /// `x`, `distances` and `labels` can be resident on the CPU or any
73
- /// GPU; copies are performed as needed
74
- void search(Index::idx_t n,
75
- const float* x,
76
- Index::idx_t k,
77
- float* distances,
78
- Index::idx_t* labels) const override;
79
-
80
- /// Overridden to force GPU indices to provide their own GPU-friendly
81
- /// implementation
82
- void compute_residual(const float* x,
83
- float* residual,
84
- Index::idx_t key) const override;
85
-
86
- /// Overridden to force GPU indices to provide their own GPU-friendly
87
- /// implementation
88
- void compute_residual_n(Index::idx_t n,
89
- const float* xs,
90
- float* residuals,
91
- const Index::idx_t* keys) const override;
92
-
93
- protected:
94
- /// Copy what we need from the CPU equivalent
95
- void copyFrom(const faiss::Index* index);
96
-
97
- /// Copy what we have to the CPU equivalent
98
- void copyTo(faiss::Index* index) const;
99
-
100
- /// Does addImpl_ require IDs? If so, and no IDs are provided, we will
101
- /// generate them sequentially based on the order in which the IDs are added
102
- virtual bool addImplRequiresIDs_() const = 0;
103
-
104
- /// Overridden to actually perform the add
105
- /// All data is guaranteed to be resident on our device
106
- virtual void addImpl_(int n,
107
- const float* x,
108
- const Index::idx_t* ids) = 0;
109
-
110
- /// Overridden to actually perform the search
111
- /// All data is guaranteed to be resident on our device
112
- virtual void searchImpl_(int n,
113
- const float* x,
114
- int k,
115
- float* distances,
116
- Index::idx_t* labels) const = 0;
117
-
118
- private:
119
- /// Handles paged adds if the add set is too large, passes to
120
- /// addImpl_ to actually perform the add for the current page
121
- void addPaged_(int n,
122
- const float* x,
123
- const Index::idx_t* ids);
124
-
125
- /// Calls addImpl_ for a single page of GPU-resident data
126
- void addPage_(int n,
127
- const float* x,
128
- const Index::idx_t* ids);
129
-
130
- /// Calls searchImpl_ for a single page of GPU-resident data
131
- void searchNonPaged_(int n,
132
- const float* x,
133
- int k,
134
- float* outDistancesData,
135
- Index::idx_t* outIndicesData) const;
136
-
137
- /// Calls searchImpl_ for a single page of GPU-resident data,
138
- /// handling paging of the data and copies from the CPU
139
- void searchFromCpuPaged_(int n,
140
- const float* x,
141
- int k,
142
- float* outDistancesData,
143
- Index::idx_t* outIndicesData) const;
144
-
145
- protected:
146
- /// Manages streams, cuBLAS handles and scratch memory for devices
147
- std::shared_ptr<GpuResources> resources_;
148
-
149
- /// Our configuration options
150
- const GpuIndexConfig config_;
151
-
152
- /// Size above which we page copies from the CPU to GPU
153
- size_t minPagedSize_;
29
+ public:
30
+ GpuIndex(
31
+ std::shared_ptr<GpuResources> resources,
32
+ int dims,
33
+ faiss::MetricType metric,
34
+ float metricArg,
35
+ GpuIndexConfig config);
36
+
37
+ /// Returns the device that this index is resident on
38
+ int getDevice() const;
39
+
40
+ /// Returns a reference to our GpuResources object that manages memory,
41
+ /// stream and handle resources on the GPU
42
+ std::shared_ptr<GpuResources> getResources();
43
+
44
+ /// Set the minimum data size for searches (in MiB) for which we use
45
+ /// CPU -> GPU paging
46
+ void setMinPagingSize(size_t size);
47
+
48
+ /// Returns the current minimum data size for paged searches
49
+ size_t getMinPagingSize() const;
50
+
51
+ /// `x` can be resident on the CPU or any GPU; copies are performed
52
+ /// as needed
53
+ /// Handles paged adds if the add set is too large; calls addInternal_
54
+ void add(Index::idx_t, const float* x) override;
55
+
56
+ /// `x` and `ids` can be resident on the CPU or any GPU; copies are
57
+ /// performed as needed
58
+ /// Handles paged adds if the add set is too large; calls addInternal_
59
+ void add_with_ids(Index::idx_t n, const float* x, const Index::idx_t* ids)
60
+ override;
61
+
62
+ /// `x` and `labels` can be resident on the CPU or any GPU; copies are
63
+ /// performed as needed
64
+ void assign(
65
+ Index::idx_t n,
66
+ const float* x,
67
+ Index::idx_t* labels,
68
+ Index::idx_t k = 1) const override;
69
+
70
+ /// `x`, `distances` and `labels` can be resident on the CPU or any
71
+ /// GPU; copies are performed as needed
72
+ void search(
73
+ Index::idx_t n,
74
+ const float* x,
75
+ Index::idx_t k,
76
+ float* distances,
77
+ Index::idx_t* labels) const override;
78
+
79
+ /// Overridden to force GPU indices to provide their own GPU-friendly
80
+ /// implementation
81
+ void compute_residual(const float* x, float* residual, Index::idx_t key)
82
+ const override;
83
+
84
+ /// Overridden to force GPU indices to provide their own GPU-friendly
85
+ /// implementation
86
+ void compute_residual_n(
87
+ Index::idx_t n,
88
+ const float* xs,
89
+ float* residuals,
90
+ const Index::idx_t* keys) const override;
91
+
92
+ protected:
93
+ /// Copy what we need from the CPU equivalent
94
+ void copyFrom(const faiss::Index* index);
95
+
96
+ /// Copy what we have to the CPU equivalent
97
+ void copyTo(faiss::Index* index) const;
98
+
99
+ /// Does addImpl_ require IDs? If so, and no IDs are provided, we will
100
+ /// generate them sequentially based on the order in which the IDs are added
101
+ virtual bool addImplRequiresIDs_() const = 0;
102
+
103
+ /// Overridden to actually perform the add
104
+ /// All data is guaranteed to be resident on our device
105
+ virtual void addImpl_(int n, const float* x, const Index::idx_t* ids) = 0;
106
+
107
+ /// Overridden to actually perform the search
108
+ /// All data is guaranteed to be resident on our device
109
+ virtual void searchImpl_(
110
+ int n,
111
+ const float* x,
112
+ int k,
113
+ float* distances,
114
+ Index::idx_t* labels) const = 0;
115
+
116
+ private:
117
+ /// Handles paged adds if the add set is too large, passes to
118
+ /// addImpl_ to actually perform the add for the current page
119
+ void addPaged_(int n, const float* x, const Index::idx_t* ids);
120
+
121
+ /// Calls addImpl_ for a single page of GPU-resident data
122
+ void addPage_(int n, const float* x, const Index::idx_t* ids);
123
+
124
+ /// Calls searchImpl_ for a single page of GPU-resident data
125
+ void searchNonPaged_(
126
+ int n,
127
+ const float* x,
128
+ int k,
129
+ float* outDistancesData,
130
+ Index::idx_t* outIndicesData) const;
131
+
132
+ /// Calls searchImpl_ for a single page of GPU-resident data,
133
+ /// handling paging of the data and copies from the CPU
134
+ void searchFromCpuPaged_(
135
+ int n,
136
+ const float* x,
137
+ int k,
138
+ float* outDistancesData,
139
+ Index::idx_t* outIndicesData) const;
140
+
141
+ protected:
142
+ /// Manages streams, cuBLAS handles and scratch memory for devices
143
+ std::shared_ptr<GpuResources> resources_;
144
+
145
+ /// Our configuration options
146
+ const GpuIndexConfig config_;
147
+
148
+ /// Size above which we page copies from the CPU to GPU
149
+ size_t minPagedSize_;
154
150
  };
155
151
 
156
- } } // namespace
152
+ } // namespace gpu
153
+ } // namespace faiss
@@ -12,85 +12,88 @@
12
12
  #include <faiss/gpu/GpuResources.h>
13
13
  #include <memory>
14
14
 
15
- namespace faiss { namespace gpu {
15
+ namespace faiss {
16
+ namespace gpu {
16
17
 
17
18
  class BinaryFlatIndex;
18
19
 
19
- struct GpuIndexBinaryFlatConfig : public GpuIndexConfig {
20
- };
20
+ struct GpuIndexBinaryFlatConfig : public GpuIndexConfig {};
21
21
 
22
22
  /// A GPU version of IndexBinaryFlat for brute-force comparison of bit vectors
23
23
  /// via Hamming distance
24
24
  class GpuIndexBinaryFlat : public IndexBinary {
25
- public:
26
- /// Construct from a pre-existing faiss::IndexBinaryFlat instance, copying
27
- /// data over to the given GPU
28
- GpuIndexBinaryFlat(GpuResourcesProvider* resources,
29
- const faiss::IndexBinaryFlat* index,
30
- GpuIndexBinaryFlatConfig config =
31
- GpuIndexBinaryFlatConfig());
32
-
33
- /// Construct an empty instance that can be added to
34
- GpuIndexBinaryFlat(GpuResourcesProvider* resources,
35
- int dims,
36
- GpuIndexBinaryFlatConfig config =
37
- GpuIndexBinaryFlatConfig());
38
-
39
- ~GpuIndexBinaryFlat() override;
40
-
41
- /// Returns the device that this index is resident on
42
- int getDevice() const;
43
-
44
- /// Returns a reference to our GpuResources object that manages memory, stream
45
- /// and handle resources on the GPU
46
- std::shared_ptr<GpuResources> getResources();
47
-
48
- /// Initialize ourselves from the given CPU index; will overwrite
49
- /// all data in ourselves
50
- void copyFrom(const faiss::IndexBinaryFlat* index);
51
-
52
- /// Copy ourselves to the given CPU index; will overwrite all data
53
- /// in the index instance
54
- void copyTo(faiss::IndexBinaryFlat* index) const;
55
-
56
- void add(faiss::IndexBinary::idx_t n,
57
- const uint8_t* x) override;
58
-
59
- void reset() override;
60
-
61
- void search(faiss::IndexBinary::idx_t n,
62
- const uint8_t* x,
63
- faiss::IndexBinary::idx_t k,
64
- int32_t* distances,
65
- faiss::IndexBinary::idx_t* labels) const override;
66
-
67
- void reconstruct(faiss::IndexBinary::idx_t key,
68
- uint8_t* recons) const override;
69
-
70
- protected:
71
- /// Called from search when the input data is on the CPU;
72
- /// potentially allows for pinned memory usage
73
- void searchFromCpuPaged_(int n,
74
- const uint8_t* x,
75
- int k,
76
- int32_t* outDistancesData,
77
- int* outIndicesData) const;
78
-
79
- void searchNonPaged_(int n,
80
- const uint8_t* x,
81
- int k,
82
- int32_t* outDistancesData,
83
- int* outIndicesData) const;
84
-
85
- protected:
86
- /// Manages streans, cuBLAS handles and scratch memory for devices
87
- std::shared_ptr<GpuResources> resources_;
88
-
89
- /// Configuration options
90
- const GpuIndexBinaryFlatConfig binaryFlatConfig_;
91
-
92
- /// Holds our GPU data containing the list of vectors
93
- std::unique_ptr<BinaryFlatIndex> data_;
25
+ public:
26
+ /// Construct from a pre-existing faiss::IndexBinaryFlat instance, copying
27
+ /// data over to the given GPU
28
+ GpuIndexBinaryFlat(
29
+ GpuResourcesProvider* resources,
30
+ const faiss::IndexBinaryFlat* index,
31
+ GpuIndexBinaryFlatConfig config = GpuIndexBinaryFlatConfig());
32
+
33
+ /// Construct an empty instance that can be added to
34
+ GpuIndexBinaryFlat(
35
+ GpuResourcesProvider* resources,
36
+ int dims,
37
+ GpuIndexBinaryFlatConfig config = GpuIndexBinaryFlatConfig());
38
+
39
+ ~GpuIndexBinaryFlat() override;
40
+
41
+ /// Returns the device that this index is resident on
42
+ int getDevice() const;
43
+
44
+ /// Returns a reference to our GpuResources object that manages memory,
45
+ /// stream and handle resources on the GPU
46
+ std::shared_ptr<GpuResources> getResources();
47
+
48
+ /// Initialize ourselves from the given CPU index; will overwrite
49
+ /// all data in ourselves
50
+ void copyFrom(const faiss::IndexBinaryFlat* index);
51
+
52
+ /// Copy ourselves to the given CPU index; will overwrite all data
53
+ /// in the index instance
54
+ void copyTo(faiss::IndexBinaryFlat* index) const;
55
+
56
+ void add(faiss::IndexBinary::idx_t n, const uint8_t* x) override;
57
+
58
+ void reset() override;
59
+
60
+ void search(
61
+ faiss::IndexBinary::idx_t n,
62
+ const uint8_t* x,
63
+ faiss::IndexBinary::idx_t k,
64
+ int32_t* distances,
65
+ faiss::IndexBinary::idx_t* labels) const override;
66
+
67
+ void reconstruct(faiss::IndexBinary::idx_t key, uint8_t* recons)
68
+ const override;
69
+
70
+ protected:
71
+ /// Called from search when the input data is on the CPU;
72
+ /// potentially allows for pinned memory usage
73
+ void searchFromCpuPaged_(
74
+ int n,
75
+ const uint8_t* x,
76
+ int k,
77
+ int32_t* outDistancesData,
78
+ int* outIndicesData) const;
79
+
80
+ void searchNonPaged_(
81
+ int n,
82
+ const uint8_t* x,
83
+ int k,
84
+ int32_t* outDistancesData,
85
+ int* outIndicesData) const;
86
+
87
+ protected:
88
+ /// Manages streans, cuBLAS handles and scratch memory for devices
89
+ std::shared_ptr<GpuResources> resources_;
90
+
91
+ /// Configuration options
92
+ const GpuIndexBinaryFlatConfig binaryFlatConfig_;
93
+
94
+ /// Holds our GPU data containing the list of vectors
95
+ std::unique_ptr<BinaryFlatIndex> data_;
94
96
  };
95
97
 
96
- } } // namespace gpu
98
+ } // namespace gpu
99
+ } // namespace faiss