faiss 0.1.7 → 0.2.3

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 +18 -0
  3. data/README.md +7 -7
  4. data/ext/faiss/ext.cpp +1 -1
  5. data/ext/faiss/extconf.rb +8 -2
  6. data/ext/faiss/index.cpp +102 -69
  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/version.rb +1 -1
  15. data/lib/faiss.rb +0 -5
  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 +26 -12
  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
@@ -12,40 +12,46 @@
12
12
 
13
13
  // Utilities for bit packing and unpacking CPU non-interleaved and GPU
14
14
  // interleaved by 32 encodings
15
- namespace faiss { namespace gpu {
15
+ namespace faiss {
16
+ namespace gpu {
16
17
 
17
18
  // Unpacks arbitrary bitwidth codes to a whole number of bytes per code
18
19
  // The layout of the input is (v0 d0)(v0 d1) ... (v0 dD)(v1 d0) ...
19
20
  // (bit packed)
20
21
  // The layout of the output is the same (byte packed to roundUp(bitsPerCode, 8)
21
22
  // / 8 bytes)
22
- std::vector<uint8_t> unpackNonInterleaved(std::vector<uint8_t> data,
23
- int numVecs,
24
- int dims,
25
- int bitsPerCode);
23
+ std::vector<uint8_t> unpackNonInterleaved(
24
+ std::vector<uint8_t> data,
25
+ int numVecs,
26
+ int dims,
27
+ int bitsPerCode);
26
28
 
27
29
  // Unpacks arbitrary bitwidth codes to a whole number of bytes per scalar code
28
30
  // The layout of the input is (v0 d0)(v1 d0) ... (v31 d0)(v0 d1) ...
29
31
  // (bit packed)
30
32
  // The layout of the input is (v0 d0)(v0 d1) ... (v0 dD)(v1 d0) ...
31
33
  // (byte packed)
32
- std::vector<uint8_t> unpackInterleaved(std::vector<uint8_t> data,
33
- int numVecs,
34
- int dims,
35
- int bitsPerCode);
34
+ std::vector<uint8_t> unpackInterleaved(
35
+ std::vector<uint8_t> data,
36
+ int numVecs,
37
+ int dims,
38
+ int bitsPerCode);
36
39
 
37
40
  // Packs data in the byte packed non-interleaved form to bit packed
38
41
  // non-interleaved form
39
- std::vector<uint8_t> packNonInterleaved(std::vector<uint8_t> data,
40
- int numVecs,
41
- int dims,
42
- int bitsPerCode);
42
+ std::vector<uint8_t> packNonInterleaved(
43
+ std::vector<uint8_t> data,
44
+ int numVecs,
45
+ int dims,
46
+ int bitsPerCode);
43
47
 
44
48
  // Packs data in the byte packed non-interleaved form to bit packed
45
49
  // interleaved form
46
- std::vector<uint8_t> packInterleaved(std::vector<uint8_t> data,
47
- int numVecs,
48
- int dims,
49
- int bitsPerCode);
50
-
51
- } } // namespace
50
+ std::vector<uint8_t> packInterleaved(
51
+ std::vector<uint8_t> data,
52
+ int numVecs,
53
+ int dims,
54
+ int bitsPerCode);
55
+
56
+ } // namespace gpu
57
+ } // namespace faiss
@@ -5,39 +5,41 @@
5
5
  * LICENSE file in the root directory of this source tree.
6
6
  */
7
7
 
8
-
9
8
  #include <faiss/gpu/impl/RemapIndices.h>
10
9
  #include <faiss/impl/FaissAssert.h>
11
10
 
12
- namespace faiss { namespace gpu {
11
+ namespace faiss {
12
+ namespace gpu {
13
13
 
14
14
  // Utility function to translate (list id, offset) to a user index on
15
15
  // the CPU. In a cpp in order to use OpenMP
16
16
  void ivfOffsetToUserIndex(
17
- Index::idx_t* indices,
18
- int numLists,
19
- int queries,
20
- int k,
21
- const std::vector<std::vector<Index::idx_t>>& listOffsetToUserIndex) {
22
- FAISS_ASSERT(numLists == listOffsetToUserIndex.size());
17
+ Index::idx_t* indices,
18
+ int numLists,
19
+ int queries,
20
+ int k,
21
+ const std::vector<std::vector<Index::idx_t>>& listOffsetToUserIndex) {
22
+ FAISS_ASSERT(numLists == listOffsetToUserIndex.size());
23
23
 
24
24
  #pragma omp parallel for
25
- for (int q = 0; q < queries; ++q) {
26
- for (int r = 0; r < k; ++r) {
27
- auto offsetIndex = indices[q * k + r];
25
+ for (int q = 0; q < queries; ++q) {
26
+ for (int r = 0; r < k; ++r) {
27
+ auto offsetIndex = indices[q * k + r];
28
28
 
29
- if (offsetIndex < 0) continue;
29
+ if (offsetIndex < 0)
30
+ continue;
30
31
 
31
- int listId = (int) (offsetIndex >> 32);
32
- int listOffset = (int) (offsetIndex & 0xffffffff);
32
+ int listId = (int)(offsetIndex >> 32);
33
+ int listOffset = (int)(offsetIndex & 0xffffffff);
33
34
 
34
- FAISS_ASSERT(listId < numLists);
35
- auto& listIndices = listOffsetToUserIndex[listId];
35
+ FAISS_ASSERT(listId < numLists);
36
+ auto& listIndices = listOffsetToUserIndex[listId];
36
37
 
37
- FAISS_ASSERT(listOffset < listIndices.size());
38
- indices[q * k + r] = listIndices[listOffset];
38
+ FAISS_ASSERT(listOffset < listIndices.size());
39
+ indices[q * k + r] = listIndices[listOffset];
40
+ }
39
41
  }
40
- }
41
42
  }
42
43
 
43
- } } // namespace
44
+ } // namespace gpu
45
+ } // namespace faiss
@@ -5,21 +5,22 @@
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 <vector>
13
12
 
14
- namespace faiss { namespace gpu {
13
+ namespace faiss {
14
+ namespace gpu {
15
15
 
16
16
  /// Utility function to translate (list id, offset) to a user index on
17
17
  /// the CPU. In a cpp in order to use OpenMP.
18
18
  void ivfOffsetToUserIndex(
19
- Index::idx_t* indices,
20
- int numLists,
21
- int queries,
22
- int k,
23
- const std::vector<std::vector<Index::idx_t>>& listOffsetToUserIndex);
19
+ Index::idx_t* indices,
20
+ int numLists,
21
+ int queries,
22
+ int k,
23
+ const std::vector<std::vector<Index::idx_t>>& listOffsetToUserIndex);
24
24
 
25
- } } // namespace
25
+ } // namespace gpu
26
+ } // namespace faiss
@@ -5,67 +5,62 @@
5
5
  * LICENSE file in the root directory of this source tree.
6
6
  */
7
7
 
8
-
9
8
  #include <faiss/impl/FaissAssert.h>
10
9
 
11
- namespace faiss { namespace gpu {
10
+ namespace faiss {
11
+ namespace gpu {
12
12
 
13
13
  template <typename GpuIndex>
14
14
  IndexWrapper<GpuIndex>::IndexWrapper(
15
- int numGpus,
16
- std::function<std::unique_ptr<GpuIndex>(GpuResourcesProvider*, int)> init) {
17
- FAISS_ASSERT(numGpus <= faiss::gpu::getNumDevices());
18
- for (int i = 0; i < numGpus; ++i) {
19
- auto res = std::unique_ptr<faiss::gpu::StandardGpuResources>(
20
- new StandardGpuResources);
15
+ int numGpus,
16
+ std::function<std::unique_ptr<GpuIndex>(GpuResourcesProvider*, int)>
17
+ init) {
18
+ FAISS_ASSERT(numGpus <= faiss::gpu::getNumDevices());
19
+ for (int i = 0; i < numGpus; ++i) {
20
+ auto res = std::unique_ptr<faiss::gpu::StandardGpuResources>(
21
+ new StandardGpuResources);
21
22
 
22
- subIndex.emplace_back(init(res.get(), i));
23
- resources.emplace_back(std::move(res));
24
- }
23
+ subIndex.emplace_back(init(res.get(), i));
24
+ resources.emplace_back(std::move(res));
25
+ }
25
26
 
26
- if (numGpus > 1) {
27
- // create proxy
28
- replicaIndex =
29
- std::unique_ptr<faiss::IndexReplicas>(new faiss::IndexReplicas);
27
+ if (numGpus > 1) {
28
+ // create proxy
29
+ replicaIndex =
30
+ std::unique_ptr<faiss::IndexReplicas>(new faiss::IndexReplicas);
30
31
 
31
- for (auto& index : subIndex) {
32
- replicaIndex->addIndex(index.get());
32
+ for (auto& index : subIndex) {
33
+ replicaIndex->addIndex(index.get());
34
+ }
33
35
  }
34
- }
35
36
  }
36
37
 
37
38
  template <typename GpuIndex>
38
- faiss::Index*
39
- IndexWrapper<GpuIndex>::getIndex() {
40
- if ((bool) replicaIndex) {
41
- return replicaIndex.get();
42
- } else {
43
- FAISS_ASSERT(!subIndex.empty());
44
- return subIndex.front().get();
45
- }
39
+ faiss::Index* IndexWrapper<GpuIndex>::getIndex() {
40
+ if ((bool)replicaIndex) {
41
+ return replicaIndex.get();
42
+ } else {
43
+ FAISS_ASSERT(!subIndex.empty());
44
+ return subIndex.front().get();
45
+ }
46
46
  }
47
47
 
48
48
  template <typename GpuIndex>
49
- void
50
- IndexWrapper<GpuIndex>::runOnIndices(std::function<void(GpuIndex*)> f) {
51
-
52
- if ((bool) replicaIndex) {
53
- replicaIndex->runOnIndex(
54
- [f](int, faiss::Index* index) {
55
- f(dynamic_cast<GpuIndex*>(index));
56
- });
57
- } else {
58
- FAISS_ASSERT(!subIndex.empty());
59
- f(subIndex.front().get());
60
- }
49
+ void IndexWrapper<GpuIndex>::runOnIndices(std::function<void(GpuIndex*)> f) {
50
+ if ((bool)replicaIndex) {
51
+ replicaIndex->runOnIndex([f](int, faiss::Index* index) {
52
+ f(dynamic_cast<GpuIndex*>(index));
53
+ });
54
+ } else {
55
+ FAISS_ASSERT(!subIndex.empty());
56
+ f(subIndex.front().get());
57
+ }
61
58
  }
62
59
 
63
60
  template <typename GpuIndex>
64
- void
65
- IndexWrapper<GpuIndex>::setNumProbes(int nprobe) {
66
- runOnIndices([nprobe](GpuIndex* index) {
67
- index->setNumProbes(nprobe);
68
- });
61
+ void IndexWrapper<GpuIndex>::setNumProbes(int nprobe) {
62
+ runOnIndices([nprobe](GpuIndex* index) { index->setNumProbes(nprobe); });
69
63
  }
70
64
 
71
- } }
65
+ } // namespace gpu
66
+ } // namespace faiss
@@ -5,7 +5,6 @@
5
5
  * LICENSE file in the root directory of this source tree.
6
6
  */
7
7
 
8
-
9
8
  #pragma once
10
9
 
11
10
  #include <faiss/IndexReplicas.h>
@@ -14,26 +13,29 @@
14
13
  #include <memory>
15
14
  #include <vector>
16
15
 
17
- namespace faiss { namespace gpu {
16
+ namespace faiss {
17
+ namespace gpu {
18
18
 
19
19
  // If we want to run multi-GPU, create a proxy to wrap the indices.
20
20
  // If we don't want multi-GPU, don't involve the proxy, so it doesn't
21
21
  // affect the timings.
22
22
  template <typename GpuIndex>
23
23
  struct IndexWrapper {
24
- std::vector<std::unique_ptr<faiss::gpu::StandardGpuResources>> resources;
25
- std::vector<std::unique_ptr<GpuIndex>> subIndex;
26
- std::unique_ptr<faiss::IndexReplicas> replicaIndex;
27
-
28
- IndexWrapper(
29
- int numGpus,
30
- std::function<std::unique_ptr<GpuIndex>(GpuResourcesProvider*, int)> init);
31
- faiss::Index* getIndex();
32
-
33
- void runOnIndices(std::function<void(GpuIndex*)> f);
34
- void setNumProbes(int nprobe);
24
+ std::vector<std::unique_ptr<faiss::gpu::StandardGpuResources>> resources;
25
+ std::vector<std::unique_ptr<GpuIndex>> subIndex;
26
+ std::unique_ptr<faiss::IndexReplicas> replicaIndex;
27
+
28
+ IndexWrapper(
29
+ int numGpus,
30
+ std::function<std::unique_ptr<GpuIndex>(GpuResourcesProvider*, int)>
31
+ init);
32
+ faiss::Index* getIndex();
33
+
34
+ void runOnIndices(std::function<void(GpuIndex*)> f);
35
+ void setNumProbes(int nprobe);
35
36
  };
36
37
 
37
- } }
38
+ } // namespace gpu
39
+ } // namespace faiss
38
40
 
39
41
  #include <faiss/gpu/perf/IndexWrapper-inl.h>
@@ -5,14 +5,13 @@
5
5
  * LICENSE file in the root directory of this source tree.
6
6
  */
7
7
 
8
-
9
- #include <faiss/utils/random.h>
10
8
  #include <faiss/Clustering.h>
11
9
  #include <faiss/gpu/GpuIndexFlat.h>
12
10
  #include <faiss/gpu/StandardGpuResources.h>
13
11
  #include <faiss/gpu/perf/IndexWrapper.h>
14
12
  #include <faiss/gpu/utils/DeviceUtils.h>
15
13
  #include <faiss/gpu/utils/Timer.h>
14
+ #include <faiss/utils/random.h>
16
15
  #include <gflags/gflags.h>
17
16
  #include <memory>
18
17
  #include <vector>
@@ -29,87 +28,94 @@ DEFINE_bool(transposed, false, "transposed vector storage");
29
28
  DEFINE_bool(verbose, false, "turn on clustering logging");
30
29
  DEFINE_int64(seed, -1, "specify random seed");
31
30
  DEFINE_int32(num_gpus, 1, "number of gpus to use");
32
- DEFINE_int64(min_paging_size, -1, "minimum size to use CPU -> GPU paged copies");
31
+ DEFINE_int64(
32
+ min_paging_size,
33
+ -1,
34
+ "minimum size to use CPU -> GPU paged copies");
33
35
  DEFINE_int64(pinned_mem, -1, "pinned memory allocation to use");
34
36
  DEFINE_int32(max_points, -1, "max points per centroid");
35
37
 
36
38
  using namespace faiss::gpu;
37
39
 
38
40
  int main(int argc, char** argv) {
39
- gflags::ParseCommandLineFlags(&argc, &argv, true);
40
-
41
- cudaProfilerStop();
42
-
43
- auto seed = FLAGS_seed != -1L ? FLAGS_seed : time(nullptr);
44
- printf("using seed %ld\n", seed);
45
-
46
- std::vector<float> vecs((size_t) FLAGS_num * FLAGS_dim);
47
- faiss::float_rand(vecs.data(), vecs.size(), seed);
48
-
49
- printf("K-means metric %s dim %d centroids %d num train %d niter %d\n",
50
- FLAGS_L2_metric ? "L2" : "IP",
51
- FLAGS_dim, FLAGS_k, FLAGS_num, FLAGS_niter);
52
- printf("float16 math %s\n", FLAGS_use_float16 ? "enabled" : "disabled");
53
- printf("transposed storage %s\n", FLAGS_transposed ? "enabled" : "disabled");
54
- printf("verbose %s\n", FLAGS_verbose ? "enabled" : "disabled");
55
-
56
- auto initFn = [](faiss::gpu::GpuResourcesProvider* res, int dev) ->
57
- std::unique_ptr<faiss::gpu::GpuIndexFlat> {
58
- if (FLAGS_pinned_mem >= 0) {
59
- ((faiss::gpu::StandardGpuResources*) res)->setPinnedMemory(
60
- FLAGS_pinned_mem);
61
- }
62
-
63
- GpuIndexFlatConfig config;
64
- config.device = dev;
65
- config.useFloat16 = FLAGS_use_float16;
66
- config.storeTransposed = FLAGS_transposed;
67
-
68
- auto p = std::unique_ptr<faiss::gpu::GpuIndexFlat>(
69
- FLAGS_L2_metric ?
70
- (faiss::gpu::GpuIndexFlat*)
71
- new faiss::gpu::GpuIndexFlatL2(res, FLAGS_dim, config) :
72
- (faiss::gpu::GpuIndexFlat*)
73
- new faiss::gpu::GpuIndexFlatIP(res, FLAGS_dim, config));
74
-
75
- if (FLAGS_min_paging_size >= 0) {
76
- p->setMinPagingSize(FLAGS_min_paging_size);
41
+ gflags::ParseCommandLineFlags(&argc, &argv, true);
42
+
43
+ cudaProfilerStop();
44
+
45
+ auto seed = FLAGS_seed != -1L ? FLAGS_seed : time(nullptr);
46
+ printf("using seed %ld\n", seed);
47
+
48
+ std::vector<float> vecs((size_t)FLAGS_num * FLAGS_dim);
49
+ faiss::float_rand(vecs.data(), vecs.size(), seed);
50
+
51
+ printf("K-means metric %s dim %d centroids %d num train %d niter %d\n",
52
+ FLAGS_L2_metric ? "L2" : "IP",
53
+ FLAGS_dim,
54
+ FLAGS_k,
55
+ FLAGS_num,
56
+ FLAGS_niter);
57
+ printf("float16 math %s\n", FLAGS_use_float16 ? "enabled" : "disabled");
58
+ printf("transposed storage %s\n",
59
+ FLAGS_transposed ? "enabled" : "disabled");
60
+ printf("verbose %s\n", FLAGS_verbose ? "enabled" : "disabled");
61
+
62
+ auto initFn = [](faiss::gpu::GpuResourcesProvider* res,
63
+ int dev) -> std::unique_ptr<faiss::gpu::GpuIndexFlat> {
64
+ if (FLAGS_pinned_mem >= 0) {
65
+ ((faiss::gpu::StandardGpuResources*)res)
66
+ ->setPinnedMemory(FLAGS_pinned_mem);
67
+ }
68
+
69
+ GpuIndexFlatConfig config;
70
+ config.device = dev;
71
+ config.useFloat16 = FLAGS_use_float16;
72
+ config.storeTransposed = FLAGS_transposed;
73
+
74
+ auto p = std::unique_ptr<faiss::gpu::GpuIndexFlat>(
75
+ FLAGS_L2_metric
76
+ ? (faiss::gpu::GpuIndexFlat*)new faiss::gpu::
77
+ GpuIndexFlatL2(res, FLAGS_dim, config)
78
+ : (faiss::gpu::GpuIndexFlat*)new faiss::gpu::
79
+ GpuIndexFlatIP(res, FLAGS_dim, config));
80
+
81
+ if (FLAGS_min_paging_size >= 0) {
82
+ p->setMinPagingSize(FLAGS_min_paging_size);
83
+ }
84
+ return p;
85
+ };
86
+
87
+ IndexWrapper<faiss::gpu::GpuIndexFlat> gpuIndex(FLAGS_num_gpus, initFn);
88
+
89
+ CUDA_VERIFY(cudaProfilerStart());
90
+ faiss::gpu::synchronizeAllDevices();
91
+
92
+ float gpuTime = 0.0f;
93
+
94
+ faiss::ClusteringParameters cp;
95
+ cp.niter = FLAGS_niter;
96
+ cp.verbose = FLAGS_verbose;
97
+
98
+ if (FLAGS_max_points > 0) {
99
+ cp.max_points_per_centroid = FLAGS_max_points;
77
100
  }
78
- return p;
79
- };
80
-
81
- IndexWrapper<faiss::gpu::GpuIndexFlat> gpuIndex(FLAGS_num_gpus, initFn);
82
-
83
- CUDA_VERIFY(cudaProfilerStart());
84
- faiss::gpu::synchronizeAllDevices();
85
101
 
86
- float gpuTime = 0.0f;
102
+ faiss::Clustering kmeans(FLAGS_dim, FLAGS_k, cp);
87
103
 
88
- faiss::ClusteringParameters cp;
89
- cp.niter = FLAGS_niter;
90
- cp.verbose = FLAGS_verbose;
104
+ // Time k-means
105
+ {
106
+ CpuTimer timer;
91
107
 
92
- if (FLAGS_max_points > 0) {
93
- cp.max_points_per_centroid = FLAGS_max_points;
94
- }
108
+ kmeans.train(FLAGS_num, vecs.data(), *(gpuIndex.getIndex()));
95
109
 
96
- faiss::Clustering kmeans(FLAGS_dim, FLAGS_k, cp);
97
-
98
- // Time k-means
99
- {
100
- CpuTimer timer;
101
-
102
- kmeans.train(FLAGS_num, vecs.data(), *(gpuIndex.getIndex()));
103
-
104
- // There is a device -> host copy above, so no need to time
105
- // additional synchronization with the GPU
106
- gpuTime = timer.elapsedMilliseconds();
107
- }
110
+ // There is a device -> host copy above, so no need to time
111
+ // additional synchronization with the GPU
112
+ gpuTime = timer.elapsedMilliseconds();
113
+ }
108
114
 
109
- CUDA_VERIFY(cudaProfilerStop());
110
- printf("k-means time %.3f ms\n", gpuTime);
115
+ CUDA_VERIFY(cudaProfilerStop());
116
+ printf("k-means time %.3f ms\n", gpuTime);
111
117
 
112
- CUDA_VERIFY(cudaDeviceSynchronize());
118
+ CUDA_VERIFY(cudaDeviceSynchronize());
113
119
 
114
- return 0;
120
+ return 0;
115
121
  }