faiss 0.2.0 → 0.2.1

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (202) hide show
  1. checksums.yaml +4 -4
  2. data/CHANGELOG.md +4 -0
  3. data/lib/faiss/version.rb +1 -1
  4. data/vendor/faiss/faiss/AutoTune.cpp +292 -291
  5. data/vendor/faiss/faiss/AutoTune.h +55 -56
  6. data/vendor/faiss/faiss/Clustering.cpp +334 -195
  7. data/vendor/faiss/faiss/Clustering.h +88 -35
  8. data/vendor/faiss/faiss/IVFlib.cpp +171 -195
  9. data/vendor/faiss/faiss/IVFlib.h +48 -51
  10. data/vendor/faiss/faiss/Index.cpp +85 -103
  11. data/vendor/faiss/faiss/Index.h +54 -48
  12. data/vendor/faiss/faiss/Index2Layer.cpp +139 -164
  13. data/vendor/faiss/faiss/Index2Layer.h +22 -22
  14. data/vendor/faiss/faiss/IndexBinary.cpp +45 -37
  15. data/vendor/faiss/faiss/IndexBinary.h +140 -132
  16. data/vendor/faiss/faiss/IndexBinaryFlat.cpp +73 -53
  17. data/vendor/faiss/faiss/IndexBinaryFlat.h +29 -24
  18. data/vendor/faiss/faiss/IndexBinaryFromFloat.cpp +46 -43
  19. data/vendor/faiss/faiss/IndexBinaryFromFloat.h +16 -15
  20. data/vendor/faiss/faiss/IndexBinaryHNSW.cpp +215 -232
  21. data/vendor/faiss/faiss/IndexBinaryHNSW.h +25 -24
  22. data/vendor/faiss/faiss/IndexBinaryHash.cpp +182 -177
  23. data/vendor/faiss/faiss/IndexBinaryHash.h +41 -34
  24. data/vendor/faiss/faiss/IndexBinaryIVF.cpp +489 -461
  25. data/vendor/faiss/faiss/IndexBinaryIVF.h +97 -68
  26. data/vendor/faiss/faiss/IndexFlat.cpp +116 -147
  27. data/vendor/faiss/faiss/IndexFlat.h +35 -46
  28. data/vendor/faiss/faiss/IndexHNSW.cpp +372 -348
  29. data/vendor/faiss/faiss/IndexHNSW.h +57 -41
  30. data/vendor/faiss/faiss/IndexIVF.cpp +474 -454
  31. data/vendor/faiss/faiss/IndexIVF.h +146 -113
  32. data/vendor/faiss/faiss/IndexIVFFlat.cpp +248 -250
  33. data/vendor/faiss/faiss/IndexIVFFlat.h +48 -51
  34. data/vendor/faiss/faiss/IndexIVFPQ.cpp +457 -516
  35. data/vendor/faiss/faiss/IndexIVFPQ.h +74 -66
  36. data/vendor/faiss/faiss/IndexIVFPQFastScan.cpp +406 -372
  37. data/vendor/faiss/faiss/IndexIVFPQFastScan.h +82 -57
  38. data/vendor/faiss/faiss/IndexIVFPQR.cpp +104 -102
  39. data/vendor/faiss/faiss/IndexIVFPQR.h +33 -28
  40. data/vendor/faiss/faiss/IndexIVFSpectralHash.cpp +125 -133
  41. data/vendor/faiss/faiss/IndexIVFSpectralHash.h +19 -21
  42. data/vendor/faiss/faiss/IndexLSH.cpp +75 -96
  43. data/vendor/faiss/faiss/IndexLSH.h +21 -26
  44. data/vendor/faiss/faiss/IndexLattice.cpp +42 -56
  45. data/vendor/faiss/faiss/IndexLattice.h +11 -16
  46. data/vendor/faiss/faiss/IndexNNDescent.cpp +231 -0
  47. data/vendor/faiss/faiss/IndexNNDescent.h +72 -0
  48. data/vendor/faiss/faiss/IndexNSG.cpp +303 -0
  49. data/vendor/faiss/faiss/IndexNSG.h +85 -0
  50. data/vendor/faiss/faiss/IndexPQ.cpp +405 -464
  51. data/vendor/faiss/faiss/IndexPQ.h +64 -67
  52. data/vendor/faiss/faiss/IndexPQFastScan.cpp +143 -170
  53. data/vendor/faiss/faiss/IndexPQFastScan.h +46 -32
  54. data/vendor/faiss/faiss/IndexPreTransform.cpp +120 -150
  55. data/vendor/faiss/faiss/IndexPreTransform.h +33 -36
  56. data/vendor/faiss/faiss/IndexRefine.cpp +115 -131
  57. data/vendor/faiss/faiss/IndexRefine.h +22 -23
  58. data/vendor/faiss/faiss/IndexReplicas.cpp +147 -153
  59. data/vendor/faiss/faiss/IndexReplicas.h +62 -56
  60. data/vendor/faiss/faiss/IndexResidual.cpp +291 -0
  61. data/vendor/faiss/faiss/IndexResidual.h +152 -0
  62. data/vendor/faiss/faiss/IndexScalarQuantizer.cpp +120 -155
  63. data/vendor/faiss/faiss/IndexScalarQuantizer.h +41 -45
  64. data/vendor/faiss/faiss/IndexShards.cpp +256 -240
  65. data/vendor/faiss/faiss/IndexShards.h +85 -73
  66. data/vendor/faiss/faiss/MatrixStats.cpp +112 -97
  67. data/vendor/faiss/faiss/MatrixStats.h +7 -10
  68. data/vendor/faiss/faiss/MetaIndexes.cpp +135 -157
  69. data/vendor/faiss/faiss/MetaIndexes.h +40 -34
  70. data/vendor/faiss/faiss/MetricType.h +7 -7
  71. data/vendor/faiss/faiss/VectorTransform.cpp +652 -474
  72. data/vendor/faiss/faiss/VectorTransform.h +61 -89
  73. data/vendor/faiss/faiss/clone_index.cpp +77 -73
  74. data/vendor/faiss/faiss/clone_index.h +4 -9
  75. data/vendor/faiss/faiss/gpu/GpuAutoTune.cpp +33 -38
  76. data/vendor/faiss/faiss/gpu/GpuAutoTune.h +11 -9
  77. data/vendor/faiss/faiss/gpu/GpuCloner.cpp +197 -170
  78. data/vendor/faiss/faiss/gpu/GpuCloner.h +53 -35
  79. data/vendor/faiss/faiss/gpu/GpuClonerOptions.cpp +12 -14
  80. data/vendor/faiss/faiss/gpu/GpuClonerOptions.h +27 -25
  81. data/vendor/faiss/faiss/gpu/GpuDistance.h +116 -112
  82. data/vendor/faiss/faiss/gpu/GpuFaissAssert.h +1 -2
  83. data/vendor/faiss/faiss/gpu/GpuIndex.h +134 -137
  84. data/vendor/faiss/faiss/gpu/GpuIndexBinaryFlat.h +76 -73
  85. data/vendor/faiss/faiss/gpu/GpuIndexFlat.h +173 -162
  86. data/vendor/faiss/faiss/gpu/GpuIndexIVF.h +67 -64
  87. data/vendor/faiss/faiss/gpu/GpuIndexIVFFlat.h +89 -86
  88. data/vendor/faiss/faiss/gpu/GpuIndexIVFPQ.h +150 -141
  89. data/vendor/faiss/faiss/gpu/GpuIndexIVFScalarQuantizer.h +101 -103
  90. data/vendor/faiss/faiss/gpu/GpuIndicesOptions.h +17 -16
  91. data/vendor/faiss/faiss/gpu/GpuResources.cpp +116 -128
  92. data/vendor/faiss/faiss/gpu/GpuResources.h +182 -186
  93. data/vendor/faiss/faiss/gpu/StandardGpuResources.cpp +433 -422
  94. data/vendor/faiss/faiss/gpu/StandardGpuResources.h +131 -130
  95. data/vendor/faiss/faiss/gpu/impl/InterleavedCodes.cpp +468 -456
  96. data/vendor/faiss/faiss/gpu/impl/InterleavedCodes.h +25 -19
  97. data/vendor/faiss/faiss/gpu/impl/RemapIndices.cpp +22 -20
  98. data/vendor/faiss/faiss/gpu/impl/RemapIndices.h +9 -8
  99. data/vendor/faiss/faiss/gpu/perf/IndexWrapper-inl.h +39 -44
  100. data/vendor/faiss/faiss/gpu/perf/IndexWrapper.h +16 -14
  101. data/vendor/faiss/faiss/gpu/perf/PerfClustering.cpp +77 -71
  102. data/vendor/faiss/faiss/gpu/perf/PerfIVFPQAdd.cpp +109 -88
  103. data/vendor/faiss/faiss/gpu/perf/WriteIndex.cpp +75 -64
  104. data/vendor/faiss/faiss/gpu/test/TestCodePacking.cpp +230 -215
  105. data/vendor/faiss/faiss/gpu/test/TestGpuIndexBinaryFlat.cpp +80 -86
  106. data/vendor/faiss/faiss/gpu/test/TestGpuIndexFlat.cpp +284 -277
  107. data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFFlat.cpp +416 -416
  108. data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFPQ.cpp +611 -517
  109. data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFScalarQuantizer.cpp +166 -164
  110. data/vendor/faiss/faiss/gpu/test/TestGpuMemoryException.cpp +61 -53
  111. data/vendor/faiss/faiss/gpu/test/TestUtils.cpp +274 -238
  112. data/vendor/faiss/faiss/gpu/test/TestUtils.h +73 -57
  113. data/vendor/faiss/faiss/gpu/test/demo_ivfpq_indexing_gpu.cpp +47 -50
  114. data/vendor/faiss/faiss/gpu/utils/DeviceUtils.h +79 -72
  115. data/vendor/faiss/faiss/gpu/utils/StackDeviceMemory.cpp +140 -146
  116. data/vendor/faiss/faiss/gpu/utils/StackDeviceMemory.h +69 -71
  117. data/vendor/faiss/faiss/gpu/utils/StaticUtils.h +21 -16
  118. data/vendor/faiss/faiss/gpu/utils/Timer.cpp +25 -29
  119. data/vendor/faiss/faiss/gpu/utils/Timer.h +30 -29
  120. data/vendor/faiss/faiss/impl/AdditiveQuantizer.cpp +270 -0
  121. data/vendor/faiss/faiss/impl/AdditiveQuantizer.h +115 -0
  122. data/vendor/faiss/faiss/impl/AuxIndexStructures.cpp +90 -120
  123. data/vendor/faiss/faiss/impl/AuxIndexStructures.h +81 -65
  124. data/vendor/faiss/faiss/impl/FaissAssert.h +73 -58
  125. data/vendor/faiss/faiss/impl/FaissException.cpp +56 -48
  126. data/vendor/faiss/faiss/impl/FaissException.h +41 -29
  127. data/vendor/faiss/faiss/impl/HNSW.cpp +595 -611
  128. data/vendor/faiss/faiss/impl/HNSW.h +179 -200
  129. data/vendor/faiss/faiss/impl/LocalSearchQuantizer.cpp +672 -0
  130. data/vendor/faiss/faiss/impl/LocalSearchQuantizer.h +172 -0
  131. data/vendor/faiss/faiss/impl/NNDescent.cpp +487 -0
  132. data/vendor/faiss/faiss/impl/NNDescent.h +154 -0
  133. data/vendor/faiss/faiss/impl/NSG.cpp +682 -0
  134. data/vendor/faiss/faiss/impl/NSG.h +199 -0
  135. data/vendor/faiss/faiss/impl/PolysemousTraining.cpp +484 -454
  136. data/vendor/faiss/faiss/impl/PolysemousTraining.h +52 -55
  137. data/vendor/faiss/faiss/impl/ProductQuantizer-inl.h +26 -47
  138. data/vendor/faiss/faiss/impl/ProductQuantizer.cpp +469 -459
  139. data/vendor/faiss/faiss/impl/ProductQuantizer.h +76 -87
  140. data/vendor/faiss/faiss/impl/ResidualQuantizer.cpp +448 -0
  141. data/vendor/faiss/faiss/impl/ResidualQuantizer.h +130 -0
  142. data/vendor/faiss/faiss/impl/ResultHandler.h +96 -132
  143. data/vendor/faiss/faiss/impl/ScalarQuantizer.cpp +648 -701
  144. data/vendor/faiss/faiss/impl/ScalarQuantizer.h +48 -46
  145. data/vendor/faiss/faiss/impl/ThreadedIndex-inl.h +129 -131
  146. data/vendor/faiss/faiss/impl/ThreadedIndex.h +61 -55
  147. data/vendor/faiss/faiss/impl/index_read.cpp +547 -479
  148. data/vendor/faiss/faiss/impl/index_write.cpp +497 -407
  149. data/vendor/faiss/faiss/impl/io.cpp +75 -94
  150. data/vendor/faiss/faiss/impl/io.h +31 -41
  151. data/vendor/faiss/faiss/impl/io_macros.h +40 -29
  152. data/vendor/faiss/faiss/impl/lattice_Zn.cpp +137 -186
  153. data/vendor/faiss/faiss/impl/lattice_Zn.h +40 -51
  154. data/vendor/faiss/faiss/impl/platform_macros.h +29 -8
  155. data/vendor/faiss/faiss/impl/pq4_fast_scan.cpp +77 -124
  156. data/vendor/faiss/faiss/impl/pq4_fast_scan.h +39 -48
  157. data/vendor/faiss/faiss/impl/pq4_fast_scan_search_1.cpp +41 -52
  158. data/vendor/faiss/faiss/impl/pq4_fast_scan_search_qbs.cpp +80 -117
  159. data/vendor/faiss/faiss/impl/simd_result_handlers.h +109 -137
  160. data/vendor/faiss/faiss/index_factory.cpp +269 -218
  161. data/vendor/faiss/faiss/index_factory.h +6 -7
  162. data/vendor/faiss/faiss/index_io.h +23 -26
  163. data/vendor/faiss/faiss/invlists/BlockInvertedLists.cpp +67 -75
  164. data/vendor/faiss/faiss/invlists/BlockInvertedLists.h +22 -24
  165. data/vendor/faiss/faiss/invlists/DirectMap.cpp +96 -112
  166. data/vendor/faiss/faiss/invlists/DirectMap.h +29 -33
  167. data/vendor/faiss/faiss/invlists/InvertedLists.cpp +307 -364
  168. data/vendor/faiss/faiss/invlists/InvertedLists.h +151 -151
  169. data/vendor/faiss/faiss/invlists/InvertedListsIOHook.cpp +29 -34
  170. data/vendor/faiss/faiss/invlists/InvertedListsIOHook.h +17 -18
  171. data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.cpp +257 -293
  172. data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.h +50 -45
  173. data/vendor/faiss/faiss/python/python_callbacks.cpp +23 -26
  174. data/vendor/faiss/faiss/python/python_callbacks.h +9 -16
  175. data/vendor/faiss/faiss/utils/AlignedTable.h +79 -44
  176. data/vendor/faiss/faiss/utils/Heap.cpp +40 -48
  177. data/vendor/faiss/faiss/utils/Heap.h +186 -209
  178. data/vendor/faiss/faiss/utils/WorkerThread.cpp +67 -76
  179. data/vendor/faiss/faiss/utils/WorkerThread.h +32 -33
  180. data/vendor/faiss/faiss/utils/distances.cpp +301 -310
  181. data/vendor/faiss/faiss/utils/distances.h +133 -118
  182. data/vendor/faiss/faiss/utils/distances_simd.cpp +456 -516
  183. data/vendor/faiss/faiss/utils/extra_distances-inl.h +117 -0
  184. data/vendor/faiss/faiss/utils/extra_distances.cpp +113 -232
  185. data/vendor/faiss/faiss/utils/extra_distances.h +30 -29
  186. data/vendor/faiss/faiss/utils/hamming-inl.h +260 -209
  187. data/vendor/faiss/faiss/utils/hamming.cpp +375 -469
  188. data/vendor/faiss/faiss/utils/hamming.h +62 -85
  189. data/vendor/faiss/faiss/utils/ordered_key_value.h +16 -18
  190. data/vendor/faiss/faiss/utils/partitioning.cpp +393 -318
  191. data/vendor/faiss/faiss/utils/partitioning.h +26 -21
  192. data/vendor/faiss/faiss/utils/quantize_lut.cpp +78 -66
  193. data/vendor/faiss/faiss/utils/quantize_lut.h +22 -20
  194. data/vendor/faiss/faiss/utils/random.cpp +39 -63
  195. data/vendor/faiss/faiss/utils/random.h +13 -16
  196. data/vendor/faiss/faiss/utils/simdlib.h +4 -2
  197. data/vendor/faiss/faiss/utils/simdlib_avx2.h +88 -85
  198. data/vendor/faiss/faiss/utils/simdlib_emulated.h +226 -165
  199. data/vendor/faiss/faiss/utils/simdlib_neon.h +832 -0
  200. data/vendor/faiss/faiss/utils/utils.cpp +304 -287
  201. data/vendor/faiss/faiss/utils/utils.h +53 -48
  202. metadata +20 -2
@@ -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