faiss 0.1.5 → 0.2.2

Sign up to get free protection for your applications and to get access to all the features.
Files changed (219) hide show
  1. checksums.yaml +4 -4
  2. data/CHANGELOG.md +24 -0
  3. data/README.md +12 -0
  4. data/ext/faiss/ext.cpp +1 -1
  5. data/ext/faiss/extconf.rb +6 -2
  6. data/ext/faiss/index.cpp +114 -43
  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.rb +0 -5
  15. data/lib/faiss/version.rb +1 -1
  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 +24 -10
  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
@@ -9,32 +9,29 @@
9
9
 
10
10
  #include <faiss/Index.h>
11
11
 
12
-
13
12
  namespace faiss {
14
13
 
15
-
16
14
  /** Index that queries in a base_index (a fast one) and refines the
17
15
  * results with an exact search, hopefully improving the results.
18
16
  */
19
- struct IndexRefine: Index {
20
-
17
+ struct IndexRefine : Index {
21
18
  /// faster index to pre-select the vectors that should be filtered
22
- Index *base_index;
19
+ Index* base_index;
23
20
 
24
21
  /// refinement index
25
- Index *refine_index;
22
+ Index* refine_index;
26
23
 
27
- bool own_fields; ///< should the base index be deallocated?
28
- bool own_refine_index; ///< same with the refinement index
24
+ bool own_fields; ///< should the base index be deallocated?
25
+ bool own_refine_index; ///< same with the refinement index
29
26
 
30
27
  /// factor between k requested in search and the k requested from
31
28
  /// the base_index (should be >= 1)
32
29
  float k_factor = 1;
33
30
 
34
- /// intitialize from empty index
35
- IndexRefine (Index *base_index, Index *refine_index);
31
+ /// initialize from empty index
32
+ IndexRefine(Index* base_index, Index* refine_index);
36
33
 
37
- IndexRefine ();
34
+ IndexRefine();
38
35
 
39
36
  void train(idx_t n, const float* x) override;
40
37
 
@@ -43,31 +40,33 @@ struct IndexRefine: Index {
43
40
  void reset() override;
44
41
 
45
42
  void search(
46
- idx_t n, const float* x, idx_t k,
47
- float* distances, idx_t* labels) const override;
43
+ idx_t n,
44
+ const float* x,
45
+ idx_t k,
46
+ float* distances,
47
+ idx_t* labels) const override;
48
48
 
49
49
  // reconstruct is routed to the refine_index
50
- void reconstruct (idx_t key, float * recons) const override;
50
+ void reconstruct(idx_t key, float* recons) const override;
51
51
 
52
52
  ~IndexRefine() override;
53
53
  };
54
54
 
55
-
56
55
  /** Version where the refinement index is an IndexFlat. It has one additional
57
56
  * constructor that takes a table of elements to add to the flat refinement
58
57
  * index */
59
- struct IndexRefineFlat: IndexRefine {
60
- explicit IndexRefineFlat (Index *base_index);
61
- IndexRefineFlat(Index *base_index, const float *xb);
58
+ struct IndexRefineFlat : IndexRefine {
59
+ explicit IndexRefineFlat(Index* base_index);
60
+ IndexRefineFlat(Index* base_index, const float* xb);
62
61
 
63
62
  IndexRefineFlat();
64
63
 
65
64
  void search(
66
- idx_t n, const float* x, idx_t k,
67
- float* distances, idx_t* labels) const override;
68
-
65
+ idx_t n,
66
+ const float* x,
67
+ idx_t k,
68
+ float* distances,
69
+ idx_t* labels) const override;
69
70
  };
70
71
 
71
-
72
-
73
72
  } // namespace faiss
@@ -14,202 +14,196 @@ namespace faiss {
14
14
 
15
15
  template <typename IndexT>
16
16
  IndexReplicasTemplate<IndexT>::IndexReplicasTemplate(bool threaded)
17
- : ThreadedIndex<IndexT>(threaded) {
18
- }
17
+ : ThreadedIndex<IndexT>(threaded) {}
19
18
 
20
19
  template <typename IndexT>
21
20
  IndexReplicasTemplate<IndexT>::IndexReplicasTemplate(idx_t d, bool threaded)
22
- : ThreadedIndex<IndexT>(d, threaded) {
23
- }
21
+ : ThreadedIndex<IndexT>(d, threaded) {}
24
22
 
25
23
  template <typename IndexT>
26
24
  IndexReplicasTemplate<IndexT>::IndexReplicasTemplate(int d, bool threaded)
27
- : ThreadedIndex<IndexT>(d, threaded) {
28
- }
25
+ : ThreadedIndex<IndexT>(d, threaded) {}
29
26
 
30
27
  template <typename IndexT>
31
- void
32
- IndexReplicasTemplate<IndexT>::onAfterAddIndex(IndexT* index) {
33
- // Make sure that the parameters are the same for all prior indices, unless
34
- // we're the first index to be added
35
- if (this->count() > 0 && this->at(0) != index) {
36
- auto existing = this->at(0);
37
-
38
- FAISS_THROW_IF_NOT_FMT(index->ntotal == existing->ntotal,
39
- "IndexReplicas: newly added index does "
40
- "not have same number of vectors as prior index; "
41
- "prior index has %" PRId64 " vectors, new index has %" PRId64,
42
- existing->ntotal, index->ntotal);
43
-
44
- FAISS_THROW_IF_NOT_MSG(index->is_trained == existing->is_trained,
45
- "IndexReplicas: newly added index does "
46
- "not have same train status as prior index");
47
-
48
- FAISS_THROW_IF_NOT_MSG(index->d == existing->d,
49
- "IndexReplicas: newly added index does "
50
- "not have same dimension as prior index");
51
- } else {
52
- syncWithSubIndexes();
53
- }
28
+ void IndexReplicasTemplate<IndexT>::onAfterAddIndex(IndexT* index) {
29
+ // Make sure that the parameters are the same for all prior indices, unless
30
+ // we're the first index to be added
31
+ if (this->count() > 0 && this->at(0) != index) {
32
+ auto existing = this->at(0);
33
+
34
+ FAISS_THROW_IF_NOT_FMT(
35
+ index->ntotal == existing->ntotal,
36
+ "IndexReplicas: newly added index does "
37
+ "not have same number of vectors as prior index; "
38
+ "prior index has %" PRId64 " vectors, new index has %" PRId64,
39
+ existing->ntotal,
40
+ index->ntotal);
41
+
42
+ FAISS_THROW_IF_NOT_MSG(
43
+ index->is_trained == existing->is_trained,
44
+ "IndexReplicas: newly added index does "
45
+ "not have same train status as prior index");
46
+
47
+ FAISS_THROW_IF_NOT_MSG(
48
+ index->d == existing->d,
49
+ "IndexReplicas: newly added index does "
50
+ "not have same dimension as prior index");
51
+ } else {
52
+ syncWithSubIndexes();
53
+ }
54
54
  }
55
55
 
56
56
  template <typename IndexT>
57
- void
58
- IndexReplicasTemplate<IndexT>::onAfterRemoveIndex(IndexT* index) {
59
- syncWithSubIndexes();
57
+ void IndexReplicasTemplate<IndexT>::onAfterRemoveIndex(IndexT* index) {
58
+ syncWithSubIndexes();
60
59
  }
61
60
 
62
61
  template <typename IndexT>
63
- void
64
- IndexReplicasTemplate<IndexT>::train(idx_t n, const component_t* x) {
65
- auto fn =
66
- [n, x](int i, IndexT* index) {
67
- if (index->verbose) {
68
- printf("begin train replica %d on %" PRId64 " points\n", i, n);
69
- }
70
-
71
- index->train(n, x);
72
-
73
- if (index->verbose) {
74
- printf("end train replica %d\n", i);
75
- }
62
+ void IndexReplicasTemplate<IndexT>::train(idx_t n, const component_t* x) {
63
+ auto fn = [n, x](int i, IndexT* index) {
64
+ if (index->verbose) {
65
+ printf("begin train replica %d on %" PRId64 " points\n", i, n);
66
+ }
67
+
68
+ index->train(n, x);
69
+
70
+ if (index->verbose) {
71
+ printf("end train replica %d\n", i);
72
+ }
76
73
  };
77
74
 
78
- this->runOnIndex(fn);
79
- syncWithSubIndexes();
75
+ this->runOnIndex(fn);
76
+ syncWithSubIndexes();
80
77
  }
81
78
 
82
79
  template <typename IndexT>
83
- void
84
- IndexReplicasTemplate<IndexT>::add(idx_t n, const component_t* x) {
85
- auto fn =
86
- [n, x](int i, IndexT* index) {
87
- if (index->verbose) {
88
- printf("begin add replica %d on %" PRId64 " points\n", i, n);
89
- }
90
-
91
- index->add(n, x);
92
-
93
- if (index->verbose) {
94
- printf("end add replica %d\n", i);
95
- }
80
+ void IndexReplicasTemplate<IndexT>::add(idx_t n, const component_t* x) {
81
+ auto fn = [n, x](int i, IndexT* index) {
82
+ if (index->verbose) {
83
+ printf("begin add replica %d on %" PRId64 " points\n", i, n);
84
+ }
85
+
86
+ index->add(n, x);
87
+
88
+ if (index->verbose) {
89
+ printf("end add replica %d\n", i);
90
+ }
96
91
  };
97
92
 
98
- this->runOnIndex(fn);
99
- syncWithSubIndexes();
93
+ this->runOnIndex(fn);
94
+ syncWithSubIndexes();
100
95
  }
101
96
 
102
97
  template <typename IndexT>
103
- void
104
- IndexReplicasTemplate<IndexT>::reconstruct(idx_t n, component_t* x) const {
105
- FAISS_THROW_IF_NOT_MSG(this->count() > 0, "no replicas in index");
98
+ void IndexReplicasTemplate<IndexT>::reconstruct(idx_t n, component_t* x) const {
99
+ FAISS_THROW_IF_NOT_MSG(this->count() > 0, "no replicas in index");
106
100
 
107
- // Just pass to the first replica
108
- this->at(0)->reconstruct(n, x);
101
+ // Just pass to the first replica
102
+ this->at(0)->reconstruct(n, x);
109
103
  }
110
104
 
111
105
  template <typename IndexT>
112
- void
113
- IndexReplicasTemplate<IndexT>::search(idx_t n,
114
- const component_t* x,
115
- idx_t k,
116
- distance_t* distances,
117
- idx_t* labels) const {
118
- FAISS_THROW_IF_NOT_MSG(this->count() > 0, "no replicas in index");
119
-
120
- if (n == 0) {
121
- return;
122
- }
123
-
124
- auto dim = this->d;
125
- size_t componentsPerVec =
126
- sizeof(component_t) == 1 ? (dim + 7) / 8 : dim;
127
-
128
- // Partition the query by the number of indices we have
129
- faiss::Index::idx_t queriesPerIndex =
130
- (faiss::Index::idx_t) (n + this->count() - 1) /
131
- (faiss::Index::idx_t) this->count();
132
- FAISS_ASSERT(n / queriesPerIndex <= this->count());
133
-
134
- auto fn =
135
- [queriesPerIndex, componentsPerVec,
136
- n, x, k, distances, labels](int i, const IndexT* index) {
137
- faiss::Index::idx_t base = (faiss::Index::idx_t) i * queriesPerIndex;
138
-
139
- if (base < n) {
140
- auto numForIndex = std::min(queriesPerIndex, n - base);
141
-
142
- if (index->verbose) {
143
- printf("begin search replica %d on %" PRId64 " points\n",
144
- i, numForIndex);
145
- }
146
-
147
- index->search(numForIndex,
148
- x + base * componentsPerVec,
149
- k,
150
- distances + base * k,
151
- labels + base * k);
152
-
153
- if (index->verbose) {
154
- printf("end search replica %d\n", i);
106
+ void IndexReplicasTemplate<IndexT>::search(
107
+ idx_t n,
108
+ const component_t* x,
109
+ idx_t k,
110
+ distance_t* distances,
111
+ idx_t* labels) const {
112
+ FAISS_THROW_IF_NOT(k > 0);
113
+
114
+ FAISS_THROW_IF_NOT_MSG(this->count() > 0, "no replicas in index");
115
+
116
+ if (n == 0) {
117
+ return;
118
+ }
119
+
120
+ auto dim = this->d;
121
+ size_t componentsPerVec = sizeof(component_t) == 1 ? (dim + 7) / 8 : dim;
122
+
123
+ // Partition the query by the number of indices we have
124
+ faiss::Index::idx_t queriesPerIndex =
125
+ (faiss::Index::idx_t)(n + this->count() - 1) /
126
+ (faiss::Index::idx_t)this->count();
127
+ FAISS_ASSERT(n / queriesPerIndex <= this->count());
128
+
129
+ auto fn = [queriesPerIndex, componentsPerVec, n, x, k, distances, labels](
130
+ int i, const IndexT* index) {
131
+ faiss::Index::idx_t base = (faiss::Index::idx_t)i * queriesPerIndex;
132
+
133
+ if (base < n) {
134
+ auto numForIndex = std::min(queriesPerIndex, n - base);
135
+
136
+ if (index->verbose) {
137
+ printf("begin search replica %d on %" PRId64 " points\n",
138
+ i,
139
+ numForIndex);
140
+ }
141
+
142
+ index->search(
143
+ numForIndex,
144
+ x + base * componentsPerVec,
145
+ k,
146
+ distances + base * k,
147
+ labels + base * k);
148
+
149
+ if (index->verbose) {
150
+ printf("end search replica %d\n", i);
151
+ }
155
152
  }
156
- }
157
153
  };
158
154
 
159
- this->runOnIndex(fn);
155
+ this->runOnIndex(fn);
160
156
  }
161
157
 
162
158
  // FIXME: assumes that nothing is currently running on the sub-indexes, which is
163
159
  // true with the normal API, but should use the runOnIndex API instead
164
160
  template <typename IndexT>
165
- void
166
- IndexReplicasTemplate<IndexT>::syncWithSubIndexes() {
167
- if (!this->count()) {
168
- this->is_trained = false;
169
- this->ntotal = 0;
170
-
171
- return;
172
- }
173
-
174
- auto firstIndex = this->at(0);
175
- this->metric_type = firstIndex->metric_type;
176
- this->is_trained = firstIndex->is_trained;
177
- this->ntotal = firstIndex->ntotal;
178
-
179
- for (int i = 1; i < this->count(); ++i) {
180
- auto index = this->at(i);
181
- FAISS_THROW_IF_NOT(this->metric_type == index->metric_type);
182
- FAISS_THROW_IF_NOT(this->d == index->d);
183
- FAISS_THROW_IF_NOT(this->is_trained == index->is_trained);
184
- FAISS_THROW_IF_NOT(this->ntotal == index->ntotal);
185
- }
161
+ void IndexReplicasTemplate<IndexT>::syncWithSubIndexes() {
162
+ if (!this->count()) {
163
+ this->is_trained = false;
164
+ this->ntotal = 0;
165
+
166
+ return;
167
+ }
168
+
169
+ auto firstIndex = this->at(0);
170
+ this->metric_type = firstIndex->metric_type;
171
+ this->is_trained = firstIndex->is_trained;
172
+ this->ntotal = firstIndex->ntotal;
173
+
174
+ for (int i = 1; i < this->count(); ++i) {
175
+ auto index = this->at(i);
176
+ FAISS_THROW_IF_NOT(this->metric_type == index->metric_type);
177
+ FAISS_THROW_IF_NOT(this->d == index->d);
178
+ FAISS_THROW_IF_NOT(this->is_trained == index->is_trained);
179
+ FAISS_THROW_IF_NOT(this->ntotal == index->ntotal);
180
+ }
186
181
  }
187
182
 
188
183
  // No metric_type for IndexBinary
189
184
  template <>
190
- void
191
- IndexReplicasTemplate<IndexBinary>::syncWithSubIndexes() {
192
- if (!this->count()) {
193
- this->is_trained = false;
194
- this->ntotal = 0;
195
-
196
- return;
197
- }
198
-
199
- auto firstIndex = this->at(0);
200
- this->is_trained = firstIndex->is_trained;
201
- this->ntotal = firstIndex->ntotal;
202
-
203
- for (int i = 1; i < this->count(); ++i) {
204
- auto index = this->at(i);
205
- FAISS_THROW_IF_NOT(this->d == index->d);
206
- FAISS_THROW_IF_NOT(this->is_trained == index->is_trained);
207
- FAISS_THROW_IF_NOT(this->ntotal == index->ntotal);
208
- }
185
+ void IndexReplicasTemplate<IndexBinary>::syncWithSubIndexes() {
186
+ if (!this->count()) {
187
+ this->is_trained = false;
188
+ this->ntotal = 0;
189
+
190
+ return;
191
+ }
192
+
193
+ auto firstIndex = this->at(0);
194
+ this->is_trained = firstIndex->is_trained;
195
+ this->ntotal = firstIndex->ntotal;
196
+
197
+ for (int i = 1; i < this->count(); ++i) {
198
+ auto index = this->at(i);
199
+ FAISS_THROW_IF_NOT(this->d == index->d);
200
+ FAISS_THROW_IF_NOT(this->is_trained == index->is_trained);
201
+ FAISS_THROW_IF_NOT(this->ntotal == index->ntotal);
202
+ }
209
203
  }
210
204
 
211
205
  // explicit instantiations
212
206
  template struct IndexReplicasTemplate<Index>;
213
207
  template struct IndexReplicasTemplate<IndexBinary>;
214
208
 
215
- } // namespace
209
+ } // namespace faiss
@@ -19,64 +19,70 @@ namespace faiss {
19
19
  /// Each index is managed by a separate CPU thread.
20
20
  template <typename IndexT>
21
21
  class IndexReplicasTemplate : public ThreadedIndex<IndexT> {
22
- public:
23
- using idx_t = typename IndexT::idx_t;
24
- using component_t = typename IndexT::component_t;
25
- using distance_t = typename IndexT::distance_t;
26
-
27
- /// The dimension that all sub-indices must share will be the dimension of the
28
- /// first sub-index added
29
- /// @param threaded do we use one thread per sub-index or do queries
30
- /// sequentially?
31
- explicit IndexReplicasTemplate(bool threaded = true);
32
-
33
- /// @param d the dimension that all sub-indices must share
34
- /// @param threaded do we use one thread per sub index or do queries
35
- /// sequentially?
36
- explicit IndexReplicasTemplate(idx_t d, bool threaded = true);
37
-
38
- /// int version due to the implicit bool conversion ambiguity of int as
39
- /// dimension
40
- explicit IndexReplicasTemplate(int d, bool threaded = true);
41
-
42
- /// Alias for addIndex()
43
- void add_replica(IndexT* index) { this->addIndex(index); }
44
-
45
- /// Alias for removeIndex()
46
- void remove_replica(IndexT* index) { this->removeIndex(index); }
47
-
48
- /// faiss::Index API
49
- /// All indices receive the same call
50
- void train(idx_t n, const component_t* x) override;
51
-
52
- /// faiss::Index API
53
- /// All indices receive the same call
54
- void add(idx_t n, const component_t* x) override;
55
-
56
- /// faiss::Index API
57
- /// Query is partitioned into a slice for each sub-index
58
- /// split by ceil(n / #indices) for our sub-indices
59
- void search(idx_t n,
60
- const component_t* x,
61
- idx_t k,
62
- distance_t* distances,
63
- idx_t* labels) const override;
64
-
65
- /// reconstructs from the first index
66
- void reconstruct(idx_t, component_t *v) const override;
67
-
68
- /// Synchronize the top-level index (IndexShards) with data in the sub-indices
69
- void syncWithSubIndexes();
70
-
71
- protected:
72
- /// Called just after an index is added
73
- void onAfterAddIndex(IndexT* index) override;
74
-
75
- /// Called just after an index is removed
76
- void onAfterRemoveIndex(IndexT* index) override;
22
+ public:
23
+ using idx_t = typename IndexT::idx_t;
24
+ using component_t = typename IndexT::component_t;
25
+ using distance_t = typename IndexT::distance_t;
26
+
27
+ /// The dimension that all sub-indices must share will be the dimension of
28
+ /// the first sub-index added
29
+ /// @param threaded do we use one thread per sub-index or do queries
30
+ /// sequentially?
31
+ explicit IndexReplicasTemplate(bool threaded = true);
32
+
33
+ /// @param d the dimension that all sub-indices must share
34
+ /// @param threaded do we use one thread per sub index or do queries
35
+ /// sequentially?
36
+ explicit IndexReplicasTemplate(idx_t d, bool threaded = true);
37
+
38
+ /// int version due to the implicit bool conversion ambiguity of int as
39
+ /// dimension
40
+ explicit IndexReplicasTemplate(int d, bool threaded = true);
41
+
42
+ /// Alias for addIndex()
43
+ void add_replica(IndexT* index) {
44
+ this->addIndex(index);
45
+ }
46
+
47
+ /// Alias for removeIndex()
48
+ void remove_replica(IndexT* index) {
49
+ this->removeIndex(index);
50
+ }
51
+
52
+ /// faiss::Index API
53
+ /// All indices receive the same call
54
+ void train(idx_t n, const component_t* x) override;
55
+
56
+ /// faiss::Index API
57
+ /// All indices receive the same call
58
+ void add(idx_t n, const component_t* x) override;
59
+
60
+ /// faiss::Index API
61
+ /// Query is partitioned into a slice for each sub-index
62
+ /// split by ceil(n / #indices) for our sub-indices
63
+ void search(
64
+ idx_t n,
65
+ const component_t* x,
66
+ idx_t k,
67
+ distance_t* distances,
68
+ idx_t* labels) const override;
69
+
70
+ /// reconstructs from the first index
71
+ void reconstruct(idx_t, component_t* v) const override;
72
+
73
+ /// Synchronize the top-level index (IndexShards) with data in the
74
+ /// sub-indices
75
+ void syncWithSubIndexes();
76
+
77
+ protected:
78
+ /// Called just after an index is added
79
+ void onAfterAddIndex(IndexT* index) override;
80
+
81
+ /// Called just after an index is removed
82
+ void onAfterRemoveIndex(IndexT* index) override;
77
83
  };
78
84
 
79
85
  using IndexReplicas = IndexReplicasTemplate<Index>;
80
86
  using IndexBinaryReplicas = IndexReplicasTemplate<IndexBinary>;
81
87
 
82
- } // namespace
88
+ } // namespace faiss