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