faiss 0.1.7 → 0.2.3

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 (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
@@ -5,51 +5,52 @@
5
5
  * LICENSE file in the root directory of this source tree.
6
6
  */
7
7
 
8
-
9
8
  #include <faiss/gpu/GpuCloner.h>
9
+ #include <faiss/impl/FaissAssert.h>
10
10
  #include <typeinfo>
11
11
 
12
- #include <faiss/gpu/GpuIndex.h>
13
- #include <faiss/impl/FaissAssert.h>
14
- #include <faiss/index_io.h>
12
+ #include <faiss/gpu/StandardGpuResources.h>
13
+
15
14
  #include <faiss/IndexFlat.h>
16
15
  #include <faiss/IndexIVF.h>
17
16
  #include <faiss/IndexIVFFlat.h>
18
- #include <faiss/IndexScalarQuantizer.h>
19
17
  #include <faiss/IndexIVFPQ.h>
20
- #include <faiss/IndexReplicas.h>
21
18
  #include <faiss/IndexPreTransform.h>
19
+ #include <faiss/IndexReplicas.h>
20
+ #include <faiss/IndexScalarQuantizer.h>
22
21
  #include <faiss/MetaIndexes.h>
22
+ #include <faiss/gpu/GpuIndex.h>
23
23
  #include <faiss/gpu/GpuIndexFlat.h>
24
24
  #include <faiss/gpu/GpuIndexIVFFlat.h>
25
25
  #include <faiss/gpu/GpuIndexIVFPQ.h>
26
26
  #include <faiss/gpu/GpuIndexIVFScalarQuantizer.h>
27
27
  #include <faiss/gpu/utils/DeviceUtils.h>
28
+ #include <faiss/impl/FaissAssert.h>
29
+ #include <faiss/index_io.h>
28
30
 
29
- namespace faiss { namespace gpu {
30
-
31
+ namespace faiss {
32
+ namespace gpu {
31
33
 
32
34
  /**********************************************************
33
35
  * Cloning to CPU
34
36
  **********************************************************/
35
37
 
36
- void ToCPUCloner::merge_index(Index *dst, Index *src, bool successive_ids)
37
- {
38
- if (auto ifl = dynamic_cast<IndexFlat *>(dst)) {
39
- auto ifl2 = dynamic_cast<const IndexFlat *>(src);
38
+ void ToCPUCloner::merge_index(Index* dst, Index* src, bool successive_ids) {
39
+ if (auto ifl = dynamic_cast<IndexFlat*>(dst)) {
40
+ auto ifl2 = dynamic_cast<const IndexFlat*>(src);
40
41
  FAISS_ASSERT(ifl2);
41
42
  FAISS_ASSERT(successive_ids);
42
43
  ifl->add(ifl2->ntotal, ifl2->xb.data());
43
- } else if(auto ifl = dynamic_cast<IndexIVFFlat *>(dst)) {
44
- auto ifl2 = dynamic_cast<IndexIVFFlat *>(src);
44
+ } else if (auto ifl = dynamic_cast<IndexIVFFlat*>(dst)) {
45
+ auto ifl2 = dynamic_cast<IndexIVFFlat*>(src);
45
46
  FAISS_ASSERT(ifl2);
46
47
  ifl->merge_from(*ifl2, successive_ids ? ifl->ntotal : 0);
47
- } else if(auto ifl = dynamic_cast<IndexIVFScalarQuantizer *>(dst)) {
48
- auto ifl2 = dynamic_cast<IndexIVFScalarQuantizer *>(src);
48
+ } else if (auto ifl = dynamic_cast<IndexIVFScalarQuantizer*>(dst)) {
49
+ auto ifl2 = dynamic_cast<IndexIVFScalarQuantizer*>(src);
49
50
  FAISS_ASSERT(ifl2);
50
51
  ifl->merge_from(*ifl2, successive_ids ? ifl->ntotal : 0);
51
- } else if(auto ifl = dynamic_cast<IndexIVFPQ *>(dst)) {
52
- auto ifl2 = dynamic_cast<IndexIVFPQ *>(src);
52
+ } else if (auto ifl = dynamic_cast<IndexIVFPQ*>(dst)) {
53
+ auto ifl2 = dynamic_cast<IndexIVFPQ*>(src);
53
54
  FAISS_ASSERT(ifl2);
54
55
  ifl->merge_from(*ifl2, successive_ids ? ifl->ntotal : 0);
55
56
  } else {
@@ -57,24 +58,22 @@ void ToCPUCloner::merge_index(Index *dst, Index *src, bool successive_ids)
57
58
  }
58
59
  }
59
60
 
60
-
61
- Index *ToCPUCloner::clone_Index(const Index *index)
62
- {
63
- if(auto ifl = dynamic_cast<const GpuIndexFlat *>(index)) {
64
- IndexFlat *res = new IndexFlat();
61
+ Index* ToCPUCloner::clone_Index(const Index* index) {
62
+ if (auto ifl = dynamic_cast<const GpuIndexFlat*>(index)) {
63
+ IndexFlat* res = new IndexFlat();
65
64
  ifl->copyTo(res);
66
65
  return res;
67
- } else if(auto ifl = dynamic_cast<const GpuIndexIVFFlat *>(index)) {
68
- IndexIVFFlat *res = new IndexIVFFlat();
66
+ } else if (auto ifl = dynamic_cast<const GpuIndexIVFFlat*>(index)) {
67
+ IndexIVFFlat* res = new IndexIVFFlat();
69
68
  ifl->copyTo(res);
70
69
  return res;
71
- } else if(auto ifl =
72
- dynamic_cast<const GpuIndexIVFScalarQuantizer *>(index)) {
73
- IndexIVFScalarQuantizer *res = new IndexIVFScalarQuantizer();
70
+ } else if (
71
+ auto ifl = dynamic_cast<const GpuIndexIVFScalarQuantizer*>(index)) {
72
+ IndexIVFScalarQuantizer* res = new IndexIVFScalarQuantizer();
74
73
  ifl->copyTo(res);
75
74
  return res;
76
- } else if(auto ipq = dynamic_cast<const GpuIndexIVFPQ *>(index)) {
77
- IndexIVFPQ *res = new IndexIVFPQ();
75
+ } else if (auto ipq = dynamic_cast<const GpuIndexIVFPQ*>(index)) {
76
+ IndexIVFPQ* res = new IndexIVFPQ();
78
77
  ipq->copyTo(res);
79
78
  return res;
80
79
 
@@ -82,17 +81,17 @@ Index *ToCPUCloner::clone_Index(const Index *index)
82
81
  // objective is to make a single component out of them
83
82
  // (inverse op of ToGpuClonerMultiple)
84
83
 
85
- } else if(auto ish = dynamic_cast<const IndexShards *>(index)) {
84
+ } else if (auto ish = dynamic_cast<const IndexShards*>(index)) {
86
85
  int nshard = ish->count();
87
86
  FAISS_ASSERT(nshard > 0);
88
- Index *res = clone_Index(ish->at(0));
89
- for(int i = 1; i < ish->count(); i++) {
90
- Index *res_i = clone_Index(ish->at(i));
87
+ Index* res = clone_Index(ish->at(0));
88
+ for (int i = 1; i < ish->count(); i++) {
89
+ Index* res_i = clone_Index(ish->at(i));
91
90
  merge_index(res, res_i, ish->successive_ids);
92
91
  delete res_i;
93
92
  }
94
93
  return res;
95
- } else if(auto ipr = dynamic_cast<const IndexReplicas *>(index)) {
94
+ } else if (auto ipr = dynamic_cast<const IndexReplicas*>(index)) {
96
95
  // just clone one of the replicas
97
96
  FAISS_ASSERT(ipr->count() > 0);
98
97
  return clone_Index(ipr->at(0));
@@ -101,81 +100,97 @@ Index *ToCPUCloner::clone_Index(const Index *index)
101
100
  }
102
101
  }
103
102
 
104
- faiss::Index * index_gpu_to_cpu(const faiss::Index *gpu_index)
105
- {
103
+ faiss::Index* index_gpu_to_cpu(const faiss::Index* gpu_index) {
106
104
  ToCPUCloner cl;
107
105
  return cl.clone_Index(gpu_index);
108
106
  }
109
107
 
110
-
111
-
112
-
113
108
  /**********************************************************
114
109
  * Cloning to 1 GPU
115
110
  **********************************************************/
116
111
 
117
- ToGpuCloner::ToGpuCloner(GpuResourcesProvider *prov, int device,
118
- const GpuClonerOptions &options):
119
- GpuClonerOptions(options), provider(prov), device(device)
120
- {}
112
+ ToGpuCloner::ToGpuCloner(
113
+ GpuResourcesProvider* prov,
114
+ int device,
115
+ const GpuClonerOptions& options)
116
+ : GpuClonerOptions(options), provider(prov), device(device) {}
121
117
 
122
- Index *ToGpuCloner::clone_Index(const Index *index)
123
- {
124
- if(auto ifl = dynamic_cast<const IndexFlat *>(index)) {
118
+ Index* ToGpuCloner::clone_Index(const Index* index) {
119
+ using idx_t = Index::idx_t;
120
+ if (auto ifl = dynamic_cast<const IndexFlat*>(index)) {
125
121
  GpuIndexFlatConfig config;
126
122
  config.device = device;
127
123
  config.useFloat16 = useFloat16;
128
124
  config.storeTransposed = storeTransposed;
129
-
130
125
  return new GpuIndexFlat(provider, ifl, config);
131
- } else if(auto ifl = dynamic_cast<const faiss::IndexIVFFlat *>(index)) {
126
+ } else if (
127
+ dynamic_cast<const IndexScalarQuantizer*>(index) &&
128
+ static_cast<const IndexScalarQuantizer*>(index)->sq.qtype ==
129
+ ScalarQuantizer::QT_fp16) {
130
+ GpuIndexFlatConfig config;
131
+ config.device = device;
132
+ config.useFloat16 = true;
133
+ GpuIndexFlat* gif = new GpuIndexFlat(
134
+ provider, index->d, index->metric_type, config);
135
+ // transfer data by blocks
136
+ idx_t bs = 1024 * 1024;
137
+ for (idx_t i0 = 0; i0 < index->ntotal; i0 += bs) {
138
+ idx_t i1 = std::min(i0 + bs, index->ntotal);
139
+ std::vector<float> buffer((i1 - i0) * index->d);
140
+ index->reconstruct_n(i0, i1 - i0, buffer.data());
141
+ gif->add(i1 - i0, buffer.data());
142
+ }
143
+ assert(gif->getNumVecs() == index->ntotal);
144
+ return gif;
145
+ } else if (auto ifl = dynamic_cast<const faiss::IndexIVFFlat*>(index)) {
132
146
  GpuIndexIVFFlatConfig config;
133
147
  config.device = device;
134
148
  config.indicesOptions = indicesOptions;
135
149
  config.flatConfig.useFloat16 = useFloat16CoarseQuantizer;
136
150
  config.flatConfig.storeTransposed = storeTransposed;
137
151
 
138
- GpuIndexIVFFlat *res =
139
- new GpuIndexIVFFlat(provider,
140
- ifl->d,
141
- ifl->nlist,
142
- ifl->metric_type,
143
- config);
144
- if(reserveVecs > 0 && ifl->ntotal == 0) {
152
+ GpuIndexIVFFlat* res = new GpuIndexIVFFlat(
153
+ provider, ifl->d, ifl->nlist, ifl->metric_type, config);
154
+ if (reserveVecs > 0 && ifl->ntotal == 0) {
145
155
  res->reserveMemory(reserveVecs);
146
156
  }
147
157
 
148
158
  res->copyFrom(ifl);
149
159
  return res;
150
- } else if(auto ifl =
151
- dynamic_cast<const faiss::IndexIVFScalarQuantizer *>(index)) {
160
+ } else if (
161
+ auto ifl = dynamic_cast<const faiss::IndexIVFScalarQuantizer*>(
162
+ index)) {
152
163
  GpuIndexIVFScalarQuantizerConfig config;
153
164
  config.device = device;
154
165
  config.indicesOptions = indicesOptions;
155
166
  config.flatConfig.useFloat16 = useFloat16CoarseQuantizer;
156
167
  config.flatConfig.storeTransposed = storeTransposed;
157
168
 
158
- GpuIndexIVFScalarQuantizer *res =
159
- new GpuIndexIVFScalarQuantizer(provider,
160
- ifl->d,
161
- ifl->nlist,
162
- ifl->sq.qtype,
163
- ifl->metric_type,
164
- ifl->by_residual,
165
- config);
166
- if(reserveVecs > 0 && ifl->ntotal == 0) {
169
+ GpuIndexIVFScalarQuantizer* res = new GpuIndexIVFScalarQuantizer(
170
+ provider,
171
+ ifl->d,
172
+ ifl->nlist,
173
+ ifl->sq.qtype,
174
+ ifl->metric_type,
175
+ ifl->by_residual,
176
+ config);
177
+ if (reserveVecs > 0 && ifl->ntotal == 0) {
167
178
  res->reserveMemory(reserveVecs);
168
179
  }
169
180
 
170
181
  res->copyFrom(ifl);
171
182
  return res;
172
- } else if(auto ipq = dynamic_cast<const faiss::IndexIVFPQ *>(index)) {
173
- if(verbose)
183
+ } else if (auto ipq = dynamic_cast<const faiss::IndexIVFPQ*>(index)) {
184
+ if (verbose) {
174
185
  printf(" IndexIVFPQ size %ld -> GpuIndexIVFPQ "
175
186
  "indicesOptions=%d "
176
187
  "usePrecomputed=%d useFloat16=%d reserveVecs=%ld\n",
177
- ipq->ntotal, indicesOptions, usePrecomputed,
178
- useFloat16, reserveVecs);
188
+ ipq->ntotal,
189
+ indicesOptions,
190
+ usePrecomputed,
191
+ useFloat16,
192
+ reserveVecs);
193
+ }
179
194
  GpuIndexIVFPQConfig config;
180
195
  config.device = device;
181
196
  config.indicesOptions = indicesOptions;
@@ -184,154 +199,144 @@ Index *ToGpuCloner::clone_Index(const Index *index)
184
199
  config.useFloat16LookupTables = useFloat16;
185
200
  config.usePrecomputedTables = usePrecomputed;
186
201
 
187
- GpuIndexIVFPQ *res = new GpuIndexIVFPQ(provider, ipq, config);
202
+ GpuIndexIVFPQ* res = new GpuIndexIVFPQ(provider, ipq, config);
188
203
 
189
- if(reserveVecs > 0 && ipq->ntotal == 0) {
204
+ if (reserveVecs > 0 && ipq->ntotal == 0) {
190
205
  res->reserveMemory(reserveVecs);
191
206
  }
192
207
 
193
208
  return res;
194
209
  } else {
210
+ // default: use CPU cloner
195
211
  return Cloner::clone_Index(index);
196
212
  }
197
213
  }
198
214
 
199
-
200
- faiss::Index * index_cpu_to_gpu(
201
- GpuResourcesProvider* provider, int device,
202
- const faiss::Index *index,
203
- const GpuClonerOptions *options)
204
- {
215
+ faiss::Index* index_cpu_to_gpu(
216
+ GpuResourcesProvider* provider,
217
+ int device,
218
+ const faiss::Index* index,
219
+ const GpuClonerOptions* options) {
205
220
  GpuClonerOptions defaults;
206
221
  ToGpuCloner cl(provider, device, options ? *options : defaults);
207
222
  return cl.clone_Index(index);
208
223
  }
209
224
 
210
-
211
225
  /**********************************************************
212
226
  * Cloning to multiple GPUs
213
227
  **********************************************************/
214
228
 
215
229
  ToGpuClonerMultiple::ToGpuClonerMultiple(
216
- std::vector<GpuResourcesProvider *> & provider,
217
- std::vector<int>& devices,
218
- const GpuMultipleClonerOptions &options):
219
- GpuMultipleClonerOptions(options)
220
- {
230
+ std::vector<GpuResourcesProvider*>& provider,
231
+ std::vector<int>& devices,
232
+ const GpuMultipleClonerOptions& options)
233
+ : GpuMultipleClonerOptions(options) {
221
234
  FAISS_ASSERT(provider.size() == devices.size());
222
- for(int i = 0; i < provider.size(); i++) {
235
+ for (int i = 0; i < provider.size(); i++) {
223
236
  sub_cloners.push_back(ToGpuCloner(provider[i], devices[i], options));
224
237
  }
225
238
  }
226
239
 
227
-
228
240
  ToGpuClonerMultiple::ToGpuClonerMultiple(
229
- const std::vector<ToGpuCloner> & sub_cloners,
230
- const GpuMultipleClonerOptions &options):
231
- GpuMultipleClonerOptions(options),
232
- sub_cloners(sub_cloners)
233
- {}
234
-
235
-
236
- void ToGpuClonerMultiple::copy_ivf_shard (
237
- const IndexIVF *index_ivf, IndexIVF *idx2,
238
- long n, long i)
239
- {
241
+ const std::vector<ToGpuCloner>& sub_cloners,
242
+ const GpuMultipleClonerOptions& options)
243
+ : GpuMultipleClonerOptions(options), sub_cloners(sub_cloners) {}
244
+
245
+ void ToGpuClonerMultiple::copy_ivf_shard(
246
+ const IndexIVF* index_ivf,
247
+ IndexIVF* idx2,
248
+ long n,
249
+ long i) {
240
250
  if (shard_type == 2) {
241
251
  long i0 = i * index_ivf->ntotal / n;
242
252
  long i1 = (i + 1) * index_ivf->ntotal / n;
243
253
 
244
- if(verbose)
245
- printf("IndexShards shard %ld indices %ld:%ld\n",
246
- i, i0, i1);
254
+ if (verbose)
255
+ printf("IndexShards shard %ld indices %ld:%ld\n", i, i0, i1);
247
256
  index_ivf->copy_subset_to(*idx2, 2, i0, i1);
248
257
  FAISS_ASSERT(idx2->ntotal == i1 - i0);
249
258
  } else if (shard_type == 1) {
250
- if(verbose)
251
- printf("IndexShards shard %ld select modulo %ld = %ld\n",
252
- i, n, i);
259
+ if (verbose)
260
+ printf("IndexShards shard %ld select modulo %ld = %ld\n", i, n, i);
253
261
  index_ivf->copy_subset_to(*idx2, 1, n, i);
254
262
  } else {
255
- FAISS_THROW_FMT ("shard_type %d not implemented", shard_type);
263
+ FAISS_THROW_FMT("shard_type %d not implemented", shard_type);
256
264
  }
257
-
258
265
  }
259
266
 
260
- Index * ToGpuClonerMultiple::clone_Index_to_shards (const Index *index)
261
- {
267
+ Index* ToGpuClonerMultiple::clone_Index_to_shards(const Index* index) {
262
268
  long n = sub_cloners.size();
263
269
 
264
- auto index_ivfpq =
265
- dynamic_cast<const faiss::IndexIVFPQ *>(index);
266
- auto index_ivfflat =
267
- dynamic_cast<const faiss::IndexIVFFlat *>(index);
270
+ auto index_ivfpq = dynamic_cast<const faiss::IndexIVFPQ*>(index);
271
+ auto index_ivfflat = dynamic_cast<const faiss::IndexIVFFlat*>(index);
268
272
  auto index_ivfsq =
269
- dynamic_cast<const faiss::IndexIVFScalarQuantizer *>(index);
270
- auto index_flat =
271
- dynamic_cast<const faiss::IndexFlat *>(index);
272
- FAISS_THROW_IF_NOT_MSG (
273
- index_ivfpq || index_ivfflat || index_flat || index_ivfsq,
274
- "IndexShards implemented only for "
275
- "IndexIVFFlat, IndexIVFScalarQuantizer, "
276
- "IndexFlat and IndexIVFPQ");
273
+ dynamic_cast<const faiss::IndexIVFScalarQuantizer*>(index);
274
+ auto index_flat = dynamic_cast<const faiss::IndexFlat*>(index);
275
+ FAISS_THROW_IF_NOT_MSG(
276
+ index_ivfpq || index_ivfflat || index_flat || index_ivfsq,
277
+ "IndexShards implemented only for "
278
+ "IndexIVFFlat, IndexIVFScalarQuantizer, "
279
+ "IndexFlat and IndexIVFPQ");
277
280
 
278
281
  std::vector<faiss::Index*> shards(n);
279
282
 
280
- for(long i = 0; i < n; i++) {
283
+ for (long i = 0; i < n; i++) {
281
284
  // make a shallow copy
282
- if(reserveVecs)
283
- sub_cloners[i].reserveVecs =
284
- (reserveVecs + n - 1) / n;
285
+ if (reserveVecs)
286
+ sub_cloners[i].reserveVecs = (reserveVecs + n - 1) / n;
285
287
 
286
288
  if (index_ivfpq) {
287
289
  faiss::IndexIVFPQ idx2(
288
- index_ivfpq->quantizer, index_ivfpq->d,
289
- index_ivfpq->nlist, index_ivfpq->code_size,
290
- index_ivfpq->pq.nbits);
290
+ index_ivfpq->quantizer,
291
+ index_ivfpq->d,
292
+ index_ivfpq->nlist,
293
+ index_ivfpq->code_size,
294
+ index_ivfpq->pq.nbits);
291
295
  idx2.metric_type = index_ivfpq->metric_type;
292
296
  idx2.pq = index_ivfpq->pq;
293
297
  idx2.nprobe = index_ivfpq->nprobe;
294
298
  idx2.use_precomputed_table = 0;
295
299
  idx2.is_trained = index->is_trained;
296
- copy_ivf_shard (index_ivfpq, &idx2, n, i);
300
+ copy_ivf_shard(index_ivfpq, &idx2, n, i);
297
301
  shards[i] = sub_cloners[i].clone_Index(&idx2);
298
302
  } else if (index_ivfflat) {
299
303
  faiss::IndexIVFFlat idx2(
300
- index_ivfflat->quantizer, index->d,
301
- index_ivfflat->nlist, index_ivfflat->metric_type);
304
+ index_ivfflat->quantizer,
305
+ index->d,
306
+ index_ivfflat->nlist,
307
+ index_ivfflat->metric_type);
302
308
  idx2.nprobe = index_ivfflat->nprobe;
303
309
  idx2.is_trained = index->is_trained;
304
- copy_ivf_shard (index_ivfflat, &idx2, n, i);
310
+ copy_ivf_shard(index_ivfflat, &idx2, n, i);
305
311
  shards[i] = sub_cloners[i].clone_Index(&idx2);
306
312
  } else if (index_ivfsq) {
307
313
  faiss::IndexIVFScalarQuantizer idx2(
308
- index_ivfsq->quantizer, index->d, index_ivfsq->nlist,
309
- index_ivfsq->sq.qtype,
310
- index_ivfsq->metric_type,
311
- index_ivfsq->by_residual);
314
+ index_ivfsq->quantizer,
315
+ index->d,
316
+ index_ivfsq->nlist,
317
+ index_ivfsq->sq.qtype,
318
+ index_ivfsq->metric_type,
319
+ index_ivfsq->by_residual);
312
320
 
313
321
  idx2.nprobe = index_ivfsq->nprobe;
314
322
  idx2.is_trained = index->is_trained;
315
323
  idx2.sq = index_ivfsq->sq;
316
- copy_ivf_shard (index_ivfsq, &idx2, n, i);
324
+ copy_ivf_shard(index_ivfsq, &idx2, n, i);
317
325
  shards[i] = sub_cloners[i].clone_Index(&idx2);
318
326
  } else if (index_flat) {
319
- faiss::IndexFlat idx2 (
320
- index->d, index->metric_type);
327
+ faiss::IndexFlat idx2(index->d, index->metric_type);
321
328
  shards[i] = sub_cloners[i].clone_Index(&idx2);
322
329
  if (index->ntotal > 0) {
323
330
  long i0 = index->ntotal * i / n;
324
331
  long i1 = index->ntotal * (i + 1) / n;
325
- shards[i]->add (i1 - i0,
326
- index_flat->xb.data() + i0 * index->d);
332
+ shards[i]->add(i1 - i0, index_flat->xb.data() + i0 * index->d);
327
333
  }
328
334
  }
329
335
  }
330
336
 
331
337
  bool successive_ids = index_flat != nullptr;
332
- faiss::IndexShards *res =
333
- new faiss::IndexShards(index->d, true,
334
- successive_ids);
338
+ faiss::IndexShards* res =
339
+ new faiss::IndexShards(index->d, true, successive_ids);
335
340
 
336
341
  for (int i = 0; i < n; i++) {
337
342
  res->add_shard(shards[i]);
@@ -341,33 +346,32 @@ Index * ToGpuClonerMultiple::clone_Index_to_shards (const Index *index)
341
346
  return res;
342
347
  }
343
348
 
344
- Index *ToGpuClonerMultiple::clone_Index(const Index *index)
345
- {
349
+ Index* ToGpuClonerMultiple::clone_Index(const Index* index) {
346
350
  long n = sub_cloners.size();
347
351
  if (n == 1)
348
352
  return sub_cloners[0].clone_Index(index);
349
353
 
350
- if(dynamic_cast<const IndexFlat *>(index) ||
351
- dynamic_cast<const faiss::IndexIVFFlat *>(index) ||
352
- dynamic_cast<const faiss::IndexIVFScalarQuantizer *>(index) ||
353
- dynamic_cast<const faiss::IndexIVFPQ *>(index)) {
354
- if(!shard) {
355
- IndexReplicas * res = new IndexReplicas();
356
- for(auto & sub_cloner: sub_cloners) {
354
+ if (dynamic_cast<const IndexFlat*>(index) ||
355
+ dynamic_cast<const faiss::IndexIVFFlat*>(index) ||
356
+ dynamic_cast<const faiss::IndexIVFScalarQuantizer*>(index) ||
357
+ dynamic_cast<const faiss::IndexIVFPQ*>(index)) {
358
+ if (!shard) {
359
+ IndexReplicas* res = new IndexReplicas();
360
+ for (auto& sub_cloner : sub_cloners) {
357
361
  res->addIndex(sub_cloner.clone_Index(index));
358
362
  }
359
363
  res->own_fields = true;
360
364
  return res;
361
365
  } else {
362
- return clone_Index_to_shards (index);
366
+ return clone_Index_to_shards(index);
363
367
  }
364
- } else if(auto miq = dynamic_cast<const MultiIndexQuantizer *>(index)) {
368
+ } else if (auto miq = dynamic_cast<const MultiIndexQuantizer*>(index)) {
365
369
  if (verbose) {
366
370
  printf("cloning MultiIndexQuantizer: "
367
371
  "will be valid only for search k=1\n");
368
372
  }
369
- const ProductQuantizer & pq = miq->pq;
370
- IndexSplitVectors *splitv = new IndexSplitVectors(pq.d, true);
373
+ const ProductQuantizer& pq = miq->pq;
374
+ IndexSplitVectors* splitv = new IndexSplitVectors(pq.d, true);
371
375
  splitv->own_fields = true;
372
376
 
373
377
  for (int m = 0; m < pq.M; m++) {
@@ -377,12 +381,13 @@ Index *ToGpuClonerMultiple::clone_Index(const Index *index)
377
381
  long i1 = pq.M <= n ? (m + 1) * n / pq.M : i0 + 1;
378
382
  std::vector<ToGpuCloner> sub_cloners_2;
379
383
  sub_cloners_2.insert(
380
- sub_cloners_2.begin(), sub_cloners.begin() + i0,
381
- sub_cloners.begin() + i1);
384
+ sub_cloners_2.begin(),
385
+ sub_cloners.begin() + i0,
386
+ sub_cloners.begin() + i1);
382
387
  ToGpuClonerMultiple cm(sub_cloners_2, *this);
383
- IndexFlatL2 idxc (pq.dsub);
384
- idxc.add (pq.ksub, pq.centroids.data() + m * pq.d * pq.ksub);
385
- Index *idx2 = cm.clone_Index(&idxc);
388
+ IndexFlatL2 idxc(pq.dsub);
389
+ idxc.add(pq.ksub, pq.centroids.data() + m * pq.d * pq.ksub);
390
+ Index* idx2 = cm.clone_Index(&idxc);
386
391
  splitv->add_sub_index(idx2);
387
392
  }
388
393
  return splitv;
@@ -391,17 +396,39 @@ Index *ToGpuClonerMultiple::clone_Index(const Index *index)
391
396
  }
392
397
  }
393
398
 
394
-
395
-
396
- faiss::Index * index_cpu_to_gpu_multiple(
397
- std::vector<GpuResourcesProvider*> & provider,
398
- std::vector<int> &devices,
399
- const faiss::Index *index,
400
- const GpuMultipleClonerOptions *options)
401
- {
399
+ faiss::Index* index_cpu_to_gpu_multiple(
400
+ std::vector<GpuResourcesProvider*>& provider,
401
+ std::vector<int>& devices,
402
+ const faiss::Index* index,
403
+ const GpuMultipleClonerOptions* options) {
402
404
  GpuMultipleClonerOptions defaults;
403
405
  ToGpuClonerMultiple cl(provider, devices, options ? *options : defaults);
404
406
  return cl.clone_Index(index);
405
407
  }
406
408
 
407
- } } // namespace
409
+ GpuProgressiveDimIndexFactory::GpuProgressiveDimIndexFactory(int ngpu) {
410
+ FAISS_THROW_IF_NOT(ngpu >= 1);
411
+ devices.resize(ngpu);
412
+ vres.resize(ngpu);
413
+
414
+ for (int i = 0; i < ngpu; i++) {
415
+ vres[i] = new StandardGpuResources();
416
+ devices[i] = i;
417
+ }
418
+ ncall = 0;
419
+ }
420
+
421
+ GpuProgressiveDimIndexFactory::~GpuProgressiveDimIndexFactory() {
422
+ for (int i = 0; i < vres.size(); i++) {
423
+ delete vres[i];
424
+ }
425
+ }
426
+
427
+ Index* GpuProgressiveDimIndexFactory::operator()(int dim) {
428
+ IndexFlatL2 index(dim);
429
+ ncall++;
430
+ return index_cpu_to_gpu_multiple(vres, devices, &index, &options);
431
+ }
432
+
433
+ } // namespace gpu
434
+ } // namespace faiss