faiss 0.2.0 → 0.2.1

Sign up to get free protection for your applications and to get access to all the features.
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
@@ -11,28 +11,23 @@
11
11
 
12
12
  #pragma once
13
13
 
14
-
15
-
16
14
  namespace faiss {
17
15
 
18
16
  struct Index;
19
17
  struct IndexIVF;
20
18
  struct VectorTransform;
21
19
 
22
-
23
20
  /* cloning functions */
24
- Index *clone_index (const Index *);
21
+ Index* clone_index(const Index*);
25
22
 
26
23
  /** Cloner class, useful to override classes with other cloning
27
24
  * functions. The cloning function above just calls
28
25
  * Cloner::clone_Index. */
29
26
  struct Cloner {
30
- virtual VectorTransform *clone_VectorTransform (const VectorTransform *);
31
- virtual Index *clone_Index (const Index *);
32
- virtual IndexIVF *clone_IndexIVF (const IndexIVF *);
27
+ virtual VectorTransform* clone_VectorTransform(const VectorTransform*);
28
+ virtual Index* clone_Index(const Index*);
29
+ virtual IndexIVF* clone_IndexIVF(const IndexIVF*);
33
30
  virtual ~Cloner() {}
34
31
  };
35
32
 
36
-
37
-
38
33
  } // namespace faiss
@@ -8,19 +8,19 @@
8
8
  #include <faiss/gpu/GpuAutoTune.h>
9
9
  #include <typeinfo>
10
10
 
11
- #include <faiss/gpu/GpuIndex.h>
11
+ #include <faiss/IndexPreTransform.h>
12
12
  #include <faiss/IndexReplicas.h>
13
13
  #include <faiss/IndexShards.h>
14
- #include <faiss/IndexPreTransform.h>
14
+ #include <faiss/gpu/GpuIndex.h>
15
15
  #include <faiss/gpu/GpuIndexFlat.h>
16
16
  #include <faiss/gpu/GpuIndexIVFFlat.h>
17
17
  #include <faiss/gpu/GpuIndexIVFPQ.h>
18
18
  #include <faiss/gpu/GpuIndexIVFScalarQuantizer.h>
19
- #include <faiss/impl/FaissAssert.h>
20
19
  #include <faiss/gpu/utils/DeviceUtils.h>
20
+ #include <faiss/impl/FaissAssert.h>
21
21
 
22
- namespace faiss { namespace gpu {
23
-
22
+ namespace faiss {
23
+ namespace gpu {
24
24
 
25
25
  using namespace ::faiss;
26
26
 
@@ -28,68 +28,63 @@ using namespace ::faiss;
28
28
  * Parameters to auto-tune on GpuIndex'es
29
29
  **********************************************************/
30
30
 
31
- #define DC(classname) auto ix = dynamic_cast<const classname *>(index)
31
+ #define DC(classname) auto ix = dynamic_cast<const classname*>(index)
32
32
 
33
-
34
- void GpuParameterSpace::initialize (const Index * index)
35
- {
36
- if (DC (IndexPreTransform)) {
33
+ void GpuParameterSpace::initialize(const Index* index) {
34
+ if (DC(IndexPreTransform)) {
37
35
  index = ix->index;
38
36
  }
39
- if (DC (IndexReplicas)) {
40
- if (ix->count() == 0) return;
37
+ if (DC(IndexReplicas)) {
38
+ if (ix->count() == 0)
39
+ return;
41
40
  index = ix->at(0);
42
41
  }
43
- if (DC (IndexShards)) {
44
- if (ix->count() == 0) return;
42
+ if (DC(IndexShards)) {
43
+ if (ix->count() == 0)
44
+ return;
45
45
  index = ix->at(0);
46
46
  }
47
- if (DC (GpuIndexIVF)) {
48
- ParameterRange & pr = add_range("nprobe");
47
+ if (DC(GpuIndexIVF)) {
48
+ ParameterRange& pr = add_range("nprobe");
49
49
  for (int i = 0; i < 12; i++) {
50
50
  size_t nprobe = 1 << i;
51
- if (nprobe >= ix->getNumLists() ||
52
- nprobe > getMaxKSelection()) break;
53
- pr.values.push_back (nprobe);
51
+ if (nprobe >= ix->getNumLists() || nprobe > getMaxKSelection())
52
+ break;
53
+ pr.values.push_back(nprobe);
54
54
  }
55
55
  }
56
56
  // not sure we should call the parent initializer
57
57
  }
58
58
 
59
-
60
-
61
59
  #undef DC
62
60
  // non-const version
63
- #define DC(classname) auto *ix = dynamic_cast<classname *>(index)
64
-
61
+ #define DC(classname) auto* ix = dynamic_cast<classname*>(index)
65
62
 
66
-
67
- void GpuParameterSpace::set_index_parameter (
68
- Index * index, const std::string & name, double val) const
69
- {
70
- if (DC (IndexReplicas)) {
63
+ void GpuParameterSpace::set_index_parameter(
64
+ Index* index,
65
+ const std::string& name,
66
+ double val) const {
67
+ if (DC(IndexReplicas)) {
71
68
  for (int i = 0; i < ix->count(); i++)
72
- set_index_parameter (ix->at(i), name, val);
69
+ set_index_parameter(ix->at(i), name, val);
73
70
  return;
74
71
  }
75
72
  if (name == "nprobe") {
76
- if (DC (GpuIndexIVF)) {
77
- ix->setNumProbes (int (val));
73
+ if (DC(GpuIndexIVF)) {
74
+ ix->setNumProbes(int(val));
78
75
  return;
79
76
  }
80
77
  }
81
78
  if (name == "use_precomputed_table") {
82
- if (DC (GpuIndexIVFPQ)) {
83
- ix->setPrecomputedCodes(bool (val));
79
+ if (DC(GpuIndexIVFPQ)) {
80
+ ix->setPrecomputedCodes(bool(val));
84
81
  return;
85
82
  }
86
83
  }
87
84
 
88
85
  // maybe normal index parameters apply?
89
- ParameterSpace::set_index_parameter (index, name, val);
86
+ ParameterSpace::set_index_parameter(index, name, val);
90
87
  }
91
88
 
92
-
93
-
94
-
95
- } } // namespace
89
+ } // namespace gpu
90
+ } // namespace faiss
@@ -7,21 +7,23 @@
7
7
 
8
8
  #pragma once
9
9
 
10
- #include <faiss/Index.h>
11
10
  #include <faiss/AutoTune.h>
11
+ #include <faiss/Index.h>
12
12
 
13
- namespace faiss { namespace gpu {
14
-
13
+ namespace faiss {
14
+ namespace gpu {
15
15
 
16
16
  /// parameter space and setters for GPU indexes
17
- struct GpuParameterSpace: faiss::ParameterSpace {
17
+ struct GpuParameterSpace : faiss::ParameterSpace {
18
18
  /// initialize with reasonable parameters for the index
19
- void initialize (const faiss::Index * index) override;
19
+ void initialize(const faiss::Index* index) override;
20
20
 
21
21
  /// set a combination of parameters on an index
22
- void set_index_parameter (
23
- faiss::Index * index, const std::string & name,
24
- double val) const override;
22
+ void set_index_parameter(
23
+ faiss::Index* index,
24
+ const std::string& name,
25
+ double val) const override;
25
26
  };
26
27
 
27
- } } // namespace
28
+ } // namespace gpu
29
+ } // namespace faiss
@@ -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