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
@@ -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
  }