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