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