faiss 0.1.5 → 0.2.2

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
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