faiss 0.2.0 → 0.2.4

Sign up to get free protection for your applications and to get access to all the features.
Files changed (215) hide show
  1. checksums.yaml +4 -4
  2. data/CHANGELOG.md +16 -0
  3. data/LICENSE.txt +1 -1
  4. data/README.md +7 -7
  5. data/ext/faiss/extconf.rb +6 -3
  6. data/ext/faiss/numo.hpp +4 -4
  7. data/ext/faiss/utils.cpp +1 -1
  8. data/ext/faiss/utils.h +1 -1
  9. data/lib/faiss/version.rb +1 -1
  10. data/vendor/faiss/faiss/AutoTune.cpp +292 -291
  11. data/vendor/faiss/faiss/AutoTune.h +55 -56
  12. data/vendor/faiss/faiss/Clustering.cpp +365 -194
  13. data/vendor/faiss/faiss/Clustering.h +102 -35
  14. data/vendor/faiss/faiss/IVFlib.cpp +171 -195
  15. data/vendor/faiss/faiss/IVFlib.h +48 -51
  16. data/vendor/faiss/faiss/Index.cpp +85 -103
  17. data/vendor/faiss/faiss/Index.h +54 -48
  18. data/vendor/faiss/faiss/Index2Layer.cpp +126 -224
  19. data/vendor/faiss/faiss/Index2Layer.h +22 -36
  20. data/vendor/faiss/faiss/IndexAdditiveQuantizer.cpp +407 -0
  21. data/vendor/faiss/faiss/IndexAdditiveQuantizer.h +195 -0
  22. data/vendor/faiss/faiss/IndexBinary.cpp +45 -37
  23. data/vendor/faiss/faiss/IndexBinary.h +140 -132
  24. data/vendor/faiss/faiss/IndexBinaryFlat.cpp +73 -53
  25. data/vendor/faiss/faiss/IndexBinaryFlat.h +29 -24
  26. data/vendor/faiss/faiss/IndexBinaryFromFloat.cpp +46 -43
  27. data/vendor/faiss/faiss/IndexBinaryFromFloat.h +16 -15
  28. data/vendor/faiss/faiss/IndexBinaryHNSW.cpp +215 -232
  29. data/vendor/faiss/faiss/IndexBinaryHNSW.h +25 -24
  30. data/vendor/faiss/faiss/IndexBinaryHash.cpp +182 -177
  31. data/vendor/faiss/faiss/IndexBinaryHash.h +41 -34
  32. data/vendor/faiss/faiss/IndexBinaryIVF.cpp +489 -461
  33. data/vendor/faiss/faiss/IndexBinaryIVF.h +97 -68
  34. data/vendor/faiss/faiss/IndexFlat.cpp +115 -176
  35. data/vendor/faiss/faiss/IndexFlat.h +42 -59
  36. data/vendor/faiss/faiss/IndexFlatCodes.cpp +67 -0
  37. data/vendor/faiss/faiss/IndexFlatCodes.h +47 -0
  38. data/vendor/faiss/faiss/IndexHNSW.cpp +372 -348
  39. data/vendor/faiss/faiss/IndexHNSW.h +57 -41
  40. data/vendor/faiss/faiss/IndexIVF.cpp +545 -453
  41. data/vendor/faiss/faiss/IndexIVF.h +169 -118
  42. data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.cpp +316 -0
  43. data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.h +121 -0
  44. data/vendor/faiss/faiss/IndexIVFFlat.cpp +247 -252
  45. data/vendor/faiss/faiss/IndexIVFFlat.h +48 -51
  46. data/vendor/faiss/faiss/IndexIVFPQ.cpp +459 -517
  47. data/vendor/faiss/faiss/IndexIVFPQ.h +75 -67
  48. data/vendor/faiss/faiss/IndexIVFPQFastScan.cpp +406 -372
  49. data/vendor/faiss/faiss/IndexIVFPQFastScan.h +82 -57
  50. data/vendor/faiss/faiss/IndexIVFPQR.cpp +104 -102
  51. data/vendor/faiss/faiss/IndexIVFPQR.h +33 -28
  52. data/vendor/faiss/faiss/IndexIVFSpectralHash.cpp +163 -150
  53. data/vendor/faiss/faiss/IndexIVFSpectralHash.h +38 -25
  54. data/vendor/faiss/faiss/IndexLSH.cpp +66 -113
  55. data/vendor/faiss/faiss/IndexLSH.h +20 -38
  56. data/vendor/faiss/faiss/IndexLattice.cpp +42 -56
  57. data/vendor/faiss/faiss/IndexLattice.h +11 -16
  58. data/vendor/faiss/faiss/IndexNNDescent.cpp +229 -0
  59. data/vendor/faiss/faiss/IndexNNDescent.h +72 -0
  60. data/vendor/faiss/faiss/IndexNSG.cpp +301 -0
  61. data/vendor/faiss/faiss/IndexNSG.h +85 -0
  62. data/vendor/faiss/faiss/IndexPQ.cpp +387 -495
  63. data/vendor/faiss/faiss/IndexPQ.h +64 -82
  64. data/vendor/faiss/faiss/IndexPQFastScan.cpp +143 -170
  65. data/vendor/faiss/faiss/IndexPQFastScan.h +46 -32
  66. data/vendor/faiss/faiss/IndexPreTransform.cpp +120 -150
  67. data/vendor/faiss/faiss/IndexPreTransform.h +33 -36
  68. data/vendor/faiss/faiss/IndexRefine.cpp +139 -127
  69. data/vendor/faiss/faiss/IndexRefine.h +32 -23
  70. data/vendor/faiss/faiss/IndexReplicas.cpp +147 -153
  71. data/vendor/faiss/faiss/IndexReplicas.h +62 -56
  72. data/vendor/faiss/faiss/IndexScalarQuantizer.cpp +111 -172
  73. data/vendor/faiss/faiss/IndexScalarQuantizer.h +41 -59
  74. data/vendor/faiss/faiss/IndexShards.cpp +256 -240
  75. data/vendor/faiss/faiss/IndexShards.h +85 -73
  76. data/vendor/faiss/faiss/MatrixStats.cpp +112 -97
  77. data/vendor/faiss/faiss/MatrixStats.h +7 -10
  78. data/vendor/faiss/faiss/MetaIndexes.cpp +135 -157
  79. data/vendor/faiss/faiss/MetaIndexes.h +40 -34
  80. data/vendor/faiss/faiss/MetricType.h +7 -7
  81. data/vendor/faiss/faiss/VectorTransform.cpp +654 -475
  82. data/vendor/faiss/faiss/VectorTransform.h +64 -89
  83. data/vendor/faiss/faiss/clone_index.cpp +78 -73
  84. data/vendor/faiss/faiss/clone_index.h +4 -9
  85. data/vendor/faiss/faiss/gpu/GpuAutoTune.cpp +33 -38
  86. data/vendor/faiss/faiss/gpu/GpuAutoTune.h +11 -9
  87. data/vendor/faiss/faiss/gpu/GpuCloner.cpp +198 -171
  88. data/vendor/faiss/faiss/gpu/GpuCloner.h +53 -35
  89. data/vendor/faiss/faiss/gpu/GpuClonerOptions.cpp +12 -14
  90. data/vendor/faiss/faiss/gpu/GpuClonerOptions.h +27 -25
  91. data/vendor/faiss/faiss/gpu/GpuDistance.h +116 -112
  92. data/vendor/faiss/faiss/gpu/GpuFaissAssert.h +1 -2
  93. data/vendor/faiss/faiss/gpu/GpuIcmEncoder.h +60 -0
  94. data/vendor/faiss/faiss/gpu/GpuIndex.h +134 -137
  95. data/vendor/faiss/faiss/gpu/GpuIndexBinaryFlat.h +76 -73
  96. data/vendor/faiss/faiss/gpu/GpuIndexFlat.h +173 -162
  97. data/vendor/faiss/faiss/gpu/GpuIndexIVF.h +67 -64
  98. data/vendor/faiss/faiss/gpu/GpuIndexIVFFlat.h +89 -86
  99. data/vendor/faiss/faiss/gpu/GpuIndexIVFPQ.h +150 -141
  100. data/vendor/faiss/faiss/gpu/GpuIndexIVFScalarQuantizer.h +101 -103
  101. data/vendor/faiss/faiss/gpu/GpuIndicesOptions.h +17 -16
  102. data/vendor/faiss/faiss/gpu/GpuResources.cpp +116 -128
  103. data/vendor/faiss/faiss/gpu/GpuResources.h +182 -186
  104. data/vendor/faiss/faiss/gpu/StandardGpuResources.cpp +433 -422
  105. data/vendor/faiss/faiss/gpu/StandardGpuResources.h +131 -130
  106. data/vendor/faiss/faiss/gpu/impl/InterleavedCodes.cpp +468 -456
  107. data/vendor/faiss/faiss/gpu/impl/InterleavedCodes.h +25 -19
  108. data/vendor/faiss/faiss/gpu/impl/RemapIndices.cpp +22 -20
  109. data/vendor/faiss/faiss/gpu/impl/RemapIndices.h +9 -8
  110. data/vendor/faiss/faiss/gpu/perf/IndexWrapper-inl.h +39 -44
  111. data/vendor/faiss/faiss/gpu/perf/IndexWrapper.h +16 -14
  112. data/vendor/faiss/faiss/gpu/perf/PerfClustering.cpp +77 -71
  113. data/vendor/faiss/faiss/gpu/perf/PerfIVFPQAdd.cpp +109 -88
  114. data/vendor/faiss/faiss/gpu/perf/WriteIndex.cpp +75 -64
  115. data/vendor/faiss/faiss/gpu/test/TestCodePacking.cpp +230 -215
  116. data/vendor/faiss/faiss/gpu/test/TestGpuIndexBinaryFlat.cpp +80 -86
  117. data/vendor/faiss/faiss/gpu/test/TestGpuIndexFlat.cpp +284 -277
  118. data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFFlat.cpp +416 -416
  119. data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFPQ.cpp +611 -517
  120. data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFScalarQuantizer.cpp +166 -164
  121. data/vendor/faiss/faiss/gpu/test/TestGpuMemoryException.cpp +61 -53
  122. data/vendor/faiss/faiss/gpu/test/TestUtils.cpp +274 -238
  123. data/vendor/faiss/faiss/gpu/test/TestUtils.h +73 -57
  124. data/vendor/faiss/faiss/gpu/test/demo_ivfpq_indexing_gpu.cpp +47 -50
  125. data/vendor/faiss/faiss/gpu/utils/DeviceUtils.h +79 -72
  126. data/vendor/faiss/faiss/gpu/utils/StackDeviceMemory.cpp +140 -146
  127. data/vendor/faiss/faiss/gpu/utils/StackDeviceMemory.h +69 -71
  128. data/vendor/faiss/faiss/gpu/utils/StaticUtils.h +21 -16
  129. data/vendor/faiss/faiss/gpu/utils/Timer.cpp +25 -29
  130. data/vendor/faiss/faiss/gpu/utils/Timer.h +30 -29
  131. data/vendor/faiss/faiss/impl/AdditiveQuantizer.cpp +503 -0
  132. data/vendor/faiss/faiss/impl/AdditiveQuantizer.h +175 -0
  133. data/vendor/faiss/faiss/impl/AuxIndexStructures.cpp +90 -120
  134. data/vendor/faiss/faiss/impl/AuxIndexStructures.h +81 -65
  135. data/vendor/faiss/faiss/impl/FaissAssert.h +73 -58
  136. data/vendor/faiss/faiss/impl/FaissException.cpp +56 -48
  137. data/vendor/faiss/faiss/impl/FaissException.h +41 -29
  138. data/vendor/faiss/faiss/impl/HNSW.cpp +606 -617
  139. data/vendor/faiss/faiss/impl/HNSW.h +179 -200
  140. data/vendor/faiss/faiss/impl/LocalSearchQuantizer.cpp +855 -0
  141. data/vendor/faiss/faiss/impl/LocalSearchQuantizer.h +244 -0
  142. data/vendor/faiss/faiss/impl/NNDescent.cpp +487 -0
  143. data/vendor/faiss/faiss/impl/NNDescent.h +154 -0
  144. data/vendor/faiss/faiss/impl/NSG.cpp +679 -0
  145. data/vendor/faiss/faiss/impl/NSG.h +199 -0
  146. data/vendor/faiss/faiss/impl/PolysemousTraining.cpp +484 -454
  147. data/vendor/faiss/faiss/impl/PolysemousTraining.h +52 -55
  148. data/vendor/faiss/faiss/impl/ProductQuantizer-inl.h +26 -47
  149. data/vendor/faiss/faiss/impl/ProductQuantizer.cpp +469 -459
  150. data/vendor/faiss/faiss/impl/ProductQuantizer.h +76 -87
  151. data/vendor/faiss/faiss/impl/ResidualQuantizer.cpp +758 -0
  152. data/vendor/faiss/faiss/impl/ResidualQuantizer.h +188 -0
  153. data/vendor/faiss/faiss/impl/ResultHandler.h +96 -132
  154. data/vendor/faiss/faiss/impl/ScalarQuantizer.cpp +647 -707
  155. data/vendor/faiss/faiss/impl/ScalarQuantizer.h +48 -46
  156. data/vendor/faiss/faiss/impl/ThreadedIndex-inl.h +129 -131
  157. data/vendor/faiss/faiss/impl/ThreadedIndex.h +61 -55
  158. data/vendor/faiss/faiss/impl/index_read.cpp +631 -480
  159. data/vendor/faiss/faiss/impl/index_write.cpp +547 -407
  160. data/vendor/faiss/faiss/impl/io.cpp +76 -95
  161. data/vendor/faiss/faiss/impl/io.h +31 -41
  162. data/vendor/faiss/faiss/impl/io_macros.h +60 -29
  163. data/vendor/faiss/faiss/impl/kmeans1d.cpp +301 -0
  164. data/vendor/faiss/faiss/impl/kmeans1d.h +48 -0
  165. data/vendor/faiss/faiss/impl/lattice_Zn.cpp +137 -186
  166. data/vendor/faiss/faiss/impl/lattice_Zn.h +40 -51
  167. data/vendor/faiss/faiss/impl/platform_macros.h +29 -8
  168. data/vendor/faiss/faiss/impl/pq4_fast_scan.cpp +77 -124
  169. data/vendor/faiss/faiss/impl/pq4_fast_scan.h +39 -48
  170. data/vendor/faiss/faiss/impl/pq4_fast_scan_search_1.cpp +41 -52
  171. data/vendor/faiss/faiss/impl/pq4_fast_scan_search_qbs.cpp +80 -117
  172. data/vendor/faiss/faiss/impl/simd_result_handlers.h +109 -137
  173. data/vendor/faiss/faiss/index_factory.cpp +619 -397
  174. data/vendor/faiss/faiss/index_factory.h +8 -6
  175. data/vendor/faiss/faiss/index_io.h +23 -26
  176. data/vendor/faiss/faiss/invlists/BlockInvertedLists.cpp +67 -75
  177. data/vendor/faiss/faiss/invlists/BlockInvertedLists.h +22 -24
  178. data/vendor/faiss/faiss/invlists/DirectMap.cpp +96 -112
  179. data/vendor/faiss/faiss/invlists/DirectMap.h +29 -33
  180. data/vendor/faiss/faiss/invlists/InvertedLists.cpp +307 -364
  181. data/vendor/faiss/faiss/invlists/InvertedLists.h +151 -151
  182. data/vendor/faiss/faiss/invlists/InvertedListsIOHook.cpp +29 -34
  183. data/vendor/faiss/faiss/invlists/InvertedListsIOHook.h +17 -18
  184. data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.cpp +257 -293
  185. data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.h +50 -45
  186. data/vendor/faiss/faiss/python/python_callbacks.cpp +23 -26
  187. data/vendor/faiss/faiss/python/python_callbacks.h +9 -16
  188. data/vendor/faiss/faiss/utils/AlignedTable.h +79 -44
  189. data/vendor/faiss/faiss/utils/Heap.cpp +40 -48
  190. data/vendor/faiss/faiss/utils/Heap.h +186 -209
  191. data/vendor/faiss/faiss/utils/WorkerThread.cpp +67 -76
  192. data/vendor/faiss/faiss/utils/WorkerThread.h +32 -33
  193. data/vendor/faiss/faiss/utils/distances.cpp +305 -312
  194. data/vendor/faiss/faiss/utils/distances.h +170 -122
  195. data/vendor/faiss/faiss/utils/distances_simd.cpp +498 -508
  196. data/vendor/faiss/faiss/utils/extra_distances-inl.h +117 -0
  197. data/vendor/faiss/faiss/utils/extra_distances.cpp +113 -232
  198. data/vendor/faiss/faiss/utils/extra_distances.h +30 -29
  199. data/vendor/faiss/faiss/utils/hamming-inl.h +260 -209
  200. data/vendor/faiss/faiss/utils/hamming.cpp +375 -469
  201. data/vendor/faiss/faiss/utils/hamming.h +62 -85
  202. data/vendor/faiss/faiss/utils/ordered_key_value.h +16 -18
  203. data/vendor/faiss/faiss/utils/partitioning.cpp +393 -318
  204. data/vendor/faiss/faiss/utils/partitioning.h +26 -21
  205. data/vendor/faiss/faiss/utils/quantize_lut.cpp +78 -66
  206. data/vendor/faiss/faiss/utils/quantize_lut.h +22 -20
  207. data/vendor/faiss/faiss/utils/random.cpp +39 -63
  208. data/vendor/faiss/faiss/utils/random.h +13 -16
  209. data/vendor/faiss/faiss/utils/simdlib.h +4 -2
  210. data/vendor/faiss/faiss/utils/simdlib_avx2.h +88 -85
  211. data/vendor/faiss/faiss/utils/simdlib_emulated.h +226 -165
  212. data/vendor/faiss/faiss/utils/simdlib_neon.h +832 -0
  213. data/vendor/faiss/faiss/utils/utils.cpp +304 -287
  214. data/vendor/faiss/faiss/utils/utils.h +54 -49
  215. metadata +29 -4
@@ -9,332 +9,337 @@
9
9
 
10
10
  #include <faiss/IndexIVFFlat.h>
11
11
 
12
+ #include <omp.h>
13
+
12
14
  #include <cinttypes>
13
15
  #include <cstdio>
14
16
 
15
17
  #include <faiss/IndexFlat.h>
16
18
 
19
+ #include <faiss/impl/AuxIndexStructures.h>
20
+ #include <faiss/impl/FaissAssert.h>
17
21
  #include <faiss/utils/distances.h>
18
22
  #include <faiss/utils/utils.h>
19
- #include <faiss/impl/FaissAssert.h>
20
- #include <faiss/impl/AuxIndexStructures.h>
21
-
22
23
 
23
24
  namespace faiss {
24
25
 
25
-
26
26
  /*****************************************
27
27
  * IndexIVFFlat implementation
28
28
  ******************************************/
29
29
 
30
- IndexIVFFlat::IndexIVFFlat (Index * quantizer,
31
- size_t d, size_t nlist, MetricType metric):
32
- IndexIVF (quantizer, d, nlist, sizeof(float) * d, metric)
33
- {
30
+ IndexIVFFlat::IndexIVFFlat(
31
+ Index* quantizer,
32
+ size_t d,
33
+ size_t nlist,
34
+ MetricType metric)
35
+ : IndexIVF(quantizer, d, nlist, sizeof(float) * d, metric) {
34
36
  code_size = sizeof(float) * d;
35
37
  }
36
38
 
39
+ void IndexIVFFlat::add_core(
40
+ idx_t n,
41
+ const float* x,
42
+ const int64_t* xids,
43
+ const int64_t* coarse_idx)
37
44
 
38
- void IndexIVFFlat::add_with_ids (idx_t n, const float * x, const idx_t *xids)
39
45
  {
40
- add_core (n, x, xids, nullptr);
41
- }
42
-
43
- void IndexIVFFlat::add_core (idx_t n, const float * x, const int64_t *xids,
44
- const int64_t *precomputed_idx)
46
+ FAISS_THROW_IF_NOT(is_trained);
47
+ FAISS_THROW_IF_NOT(coarse_idx);
48
+ assert(invlists);
49
+ direct_map.check_can_add(xids);
45
50
 
46
- {
47
- FAISS_THROW_IF_NOT (is_trained);
48
- assert (invlists);
49
- direct_map.check_can_add (xids);
50
- const int64_t * idx;
51
- ScopeDeleter<int64_t> del;
52
-
53
- if (precomputed_idx) {
54
- idx = precomputed_idx;
55
- } else {
56
- int64_t * idx0 = new int64_t [n];
57
- del.set (idx0);
58
- quantizer->assign (n, x, idx0);
59
- idx = idx0;
60
- }
61
51
  int64_t n_add = 0;
62
- for (size_t i = 0; i < n; i++) {
63
- idx_t id = xids ? xids[i] : ntotal + i;
64
- idx_t list_no = idx [i];
65
- size_t offset;
66
-
67
- if (list_no >= 0) {
68
- const float *xi = x + i * d;
69
- offset = invlists->add_entry (
70
- list_no, id, (const uint8_t*) xi);
71
- n_add++;
72
- } else {
73
- offset = 0;
52
+
53
+ DirectMapAdd dm_adder(direct_map, n, xids);
54
+
55
+ #pragma omp parallel reduction(+ : n_add)
56
+ {
57
+ int nt = omp_get_num_threads();
58
+ int rank = omp_get_thread_num();
59
+
60
+ // each thread takes care of a subset of lists
61
+ for (size_t i = 0; i < n; i++) {
62
+ idx_t list_no = coarse_idx[i];
63
+
64
+ if (list_no >= 0 && list_no % nt == rank) {
65
+ idx_t id = xids ? xids[i] : ntotal + i;
66
+ const float* xi = x + i * d;
67
+ size_t offset =
68
+ invlists->add_entry(list_no, id, (const uint8_t*)xi);
69
+ dm_adder.add(i, list_no, offset);
70
+ n_add++;
71
+ } else if (rank == 0 && list_no == -1) {
72
+ dm_adder.add(i, -1, 0);
73
+ }
74
74
  }
75
- direct_map.add_single_id (id, list_no, offset);
76
75
  }
77
76
 
78
77
  if (verbose) {
79
- printf("IndexIVFFlat::add_core: added %" PRId64 " / %" PRId64 " vectors\n",
80
- n_add, n);
78
+ printf("IndexIVFFlat::add_core: added %" PRId64 " / %" PRId64
79
+ " vectors\n",
80
+ n_add,
81
+ n);
81
82
  }
82
83
  ntotal += n;
83
84
  }
84
85
 
85
- void IndexIVFFlat::encode_vectors(idx_t n, const float* x,
86
- const idx_t * list_nos,
87
- uint8_t * codes,
88
- bool include_listnos) const
89
- {
86
+ void IndexIVFFlat::encode_vectors(
87
+ idx_t n,
88
+ const float* x,
89
+ const idx_t* list_nos,
90
+ uint8_t* codes,
91
+ bool include_listnos) const {
90
92
  if (!include_listnos) {
91
- memcpy (codes, x, code_size * n);
93
+ memcpy(codes, x, code_size * n);
92
94
  } else {
93
- size_t coarse_size = coarse_code_size ();
95
+ size_t coarse_size = coarse_code_size();
94
96
  for (size_t i = 0; i < n; i++) {
95
- int64_t list_no = list_nos [i];
96
- uint8_t *code = codes + i * (code_size + coarse_size);
97
- const float *xi = x + i * d;
97
+ int64_t list_no = list_nos[i];
98
+ uint8_t* code = codes + i * (code_size + coarse_size);
99
+ const float* xi = x + i * d;
98
100
  if (list_no >= 0) {
99
- encode_listno (list_no, code);
100
- memcpy (code + coarse_size, xi, code_size);
101
+ encode_listno(list_no, code);
102
+ memcpy(code + coarse_size, xi, code_size);
101
103
  } else {
102
- memset (code, 0, code_size + coarse_size);
104
+ memset(code, 0, code_size + coarse_size);
103
105
  }
104
-
105
106
  }
106
107
  }
107
108
  }
108
109
 
109
- void IndexIVFFlat::sa_decode (idx_t n, const uint8_t *bytes,
110
- float *x) const
111
- {
112
- size_t coarse_size = coarse_code_size ();
110
+ void IndexIVFFlat::sa_decode(idx_t n, const uint8_t* bytes, float* x) const {
111
+ size_t coarse_size = coarse_code_size();
113
112
  for (size_t i = 0; i < n; i++) {
114
- const uint8_t *code = bytes + i * (code_size + coarse_size);
115
- float *xi = x + i * d;
116
- memcpy (xi, code + coarse_size, code_size);
113
+ const uint8_t* code = bytes + i * (code_size + coarse_size);
114
+ float* xi = x + i * d;
115
+ memcpy(xi, code + coarse_size, code_size);
117
116
  }
118
117
  }
119
118
 
120
-
121
119
  namespace {
122
120
 
123
-
124
- template<MetricType metric, class C>
125
- struct IVFFlatScanner: InvertedListScanner {
121
+ template <MetricType metric, class C>
122
+ struct IVFFlatScanner : InvertedListScanner {
126
123
  size_t d;
127
- bool store_pairs;
128
124
 
129
- IVFFlatScanner(size_t d, bool store_pairs):
130
- d(d), store_pairs(store_pairs) {}
125
+ IVFFlatScanner(size_t d, bool store_pairs) : d(d) {
126
+ this->store_pairs = store_pairs;
127
+ }
131
128
 
132
- const float *xi;
133
- void set_query (const float *query) override {
129
+ const float* xi;
130
+ void set_query(const float* query) override {
134
131
  this->xi = query;
135
132
  }
136
133
 
137
- idx_t list_no;
138
- void set_list (idx_t list_no, float /* coarse_dis */) override {
134
+ void set_list(idx_t list_no, float /* coarse_dis */) override {
139
135
  this->list_no = list_no;
140
136
  }
141
137
 
142
- float distance_to_code (const uint8_t *code) const override {
143
- const float *yj = (float*)code;
144
- float dis = metric == METRIC_INNER_PRODUCT ?
145
- fvec_inner_product (xi, yj, d) : fvec_L2sqr (xi, yj, d);
138
+ float distance_to_code(const uint8_t* code) const override {
139
+ const float* yj = (float*)code;
140
+ float dis = metric == METRIC_INNER_PRODUCT
141
+ ? fvec_inner_product(xi, yj, d)
142
+ : fvec_L2sqr(xi, yj, d);
146
143
  return dis;
147
144
  }
148
145
 
149
- size_t scan_codes (size_t list_size,
150
- const uint8_t *codes,
151
- const idx_t *ids,
152
- float *simi, idx_t *idxi,
153
- size_t k) const override
154
- {
155
- const float *list_vecs = (const float*)codes;
146
+ size_t scan_codes(
147
+ size_t list_size,
148
+ const uint8_t* codes,
149
+ const idx_t* ids,
150
+ float* simi,
151
+ idx_t* idxi,
152
+ size_t k) const override {
153
+ const float* list_vecs = (const float*)codes;
156
154
  size_t nup = 0;
157
155
  for (size_t j = 0; j < list_size; j++) {
158
- const float * yj = list_vecs + d * j;
159
- float dis = metric == METRIC_INNER_PRODUCT ?
160
- fvec_inner_product (xi, yj, d) : fvec_L2sqr (xi, yj, d);
161
- if (C::cmp (simi[0], dis)) {
162
- int64_t id = store_pairs ? lo_build (list_no, j) : ids[j];
163
- heap_replace_top<C> (k, simi, idxi, dis, id);
156
+ const float* yj = list_vecs + d * j;
157
+ float dis = metric == METRIC_INNER_PRODUCT
158
+ ? fvec_inner_product(xi, yj, d)
159
+ : fvec_L2sqr(xi, yj, d);
160
+ if (C::cmp(simi[0], dis)) {
161
+ int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
162
+ heap_replace_top<C>(k, simi, idxi, dis, id);
164
163
  nup++;
165
164
  }
166
165
  }
167
166
  return nup;
168
167
  }
169
168
 
170
- void scan_codes_range (size_t list_size,
171
- const uint8_t *codes,
172
- const idx_t *ids,
173
- float radius,
174
- RangeQueryResult & res) const override
175
- {
176
- const float *list_vecs = (const float*)codes;
169
+ void scan_codes_range(
170
+ size_t list_size,
171
+ const uint8_t* codes,
172
+ const idx_t* ids,
173
+ float radius,
174
+ RangeQueryResult& res) const override {
175
+ const float* list_vecs = (const float*)codes;
177
176
  for (size_t j = 0; j < list_size; j++) {
178
- const float * yj = list_vecs + d * j;
179
- float dis = metric == METRIC_INNER_PRODUCT ?
180
- fvec_inner_product (xi, yj, d) : fvec_L2sqr (xi, yj, d);
181
- if (C::cmp (radius, dis)) {
182
- int64_t id = store_pairs ? lo_build (list_no, j) : ids[j];
183
- res.add (dis, id);
177
+ const float* yj = list_vecs + d * j;
178
+ float dis = metric == METRIC_INNER_PRODUCT
179
+ ? fvec_inner_product(xi, yj, d)
180
+ : fvec_L2sqr(xi, yj, d);
181
+ if (C::cmp(radius, dis)) {
182
+ int64_t id = store_pairs ? lo_build(list_no, j) : ids[j];
183
+ res.add(dis, id);
184
184
  }
185
185
  }
186
186
  }
187
-
188
-
189
187
  };
190
188
 
191
-
192
189
  } // anonymous namespace
193
190
 
194
-
195
-
196
- InvertedListScanner* IndexIVFFlat::get_InvertedListScanner
197
- (bool store_pairs) const
198
- {
191
+ InvertedListScanner* IndexIVFFlat::get_InvertedListScanner(
192
+ bool store_pairs) const {
199
193
  if (metric_type == METRIC_INNER_PRODUCT) {
200
- return new IVFFlatScanner<
201
- METRIC_INNER_PRODUCT, CMin<float, int64_t> > (d, store_pairs);
194
+ return new IVFFlatScanner<METRIC_INNER_PRODUCT, CMin<float, int64_t>>(
195
+ d, store_pairs);
202
196
  } else if (metric_type == METRIC_L2) {
203
- return new IVFFlatScanner<
204
- METRIC_L2, CMax<float, int64_t> >(d, store_pairs);
197
+ return new IVFFlatScanner<METRIC_L2, CMax<float, int64_t>>(
198
+ d, store_pairs);
205
199
  } else {
206
200
  FAISS_THROW_MSG("metric type not supported");
207
201
  }
208
202
  return nullptr;
209
203
  }
210
204
 
211
-
212
-
213
-
214
- void IndexIVFFlat::reconstruct_from_offset (int64_t list_no, int64_t offset,
215
- float* recons) const
216
- {
217
- memcpy (recons, invlists->get_single_code (list_no, offset), code_size);
205
+ void IndexIVFFlat::reconstruct_from_offset(
206
+ int64_t list_no,
207
+ int64_t offset,
208
+ float* recons) const {
209
+ memcpy(recons, invlists->get_single_code(list_no, offset), code_size);
218
210
  }
219
211
 
220
212
  /*****************************************
221
213
  * IndexIVFFlatDedup implementation
222
214
  ******************************************/
223
215
 
224
- IndexIVFFlatDedup::IndexIVFFlatDedup (
225
- Index * quantizer, size_t d, size_t nlist_,
226
- MetricType metric_type):
227
- IndexIVFFlat (quantizer, d, nlist_, metric_type)
228
- {}
216
+ IndexIVFFlatDedup::IndexIVFFlatDedup(
217
+ Index* quantizer,
218
+ size_t d,
219
+ size_t nlist_,
220
+ MetricType metric_type)
221
+ : IndexIVFFlat(quantizer, d, nlist_, metric_type) {}
229
222
 
230
-
231
- void IndexIVFFlatDedup::train(idx_t n, const float* x)
232
- {
223
+ void IndexIVFFlatDedup::train(idx_t n, const float* x) {
233
224
  std::unordered_map<uint64_t, idx_t> map;
234
- float * x2 = new float [n * d];
235
- ScopeDeleter<float> del (x2);
225
+ std::unique_ptr<float[]> x2(new float[n * d]);
236
226
 
237
227
  int64_t n2 = 0;
238
228
  for (int64_t i = 0; i < n; i++) {
239
- uint64_t hash = hash_bytes((uint8_t *)(x + i * d), code_size);
229
+ uint64_t hash = hash_bytes((uint8_t*)(x + i * d), code_size);
240
230
  if (map.count(hash) &&
241
- !memcmp (x2 + map[hash] * d, x + i * d, code_size)) {
231
+ !memcmp(x2.get() + map[hash] * d, x + i * d, code_size)) {
242
232
  // is duplicate, skip
243
233
  } else {
244
- map [hash] = n2;
245
- memcpy (x2 + n2 * d, x + i * d, code_size);
246
- n2 ++;
234
+ map[hash] = n2;
235
+ memcpy(x2.get() + n2 * d, x + i * d, code_size);
236
+ n2++;
247
237
  }
248
238
  }
249
239
  if (verbose) {
250
- printf ("IndexIVFFlatDedup::train: train on %" PRId64 " points after dedup "
251
- "(was %" PRId64 " points)\n", n2, n);
240
+ printf("IndexIVFFlatDedup::train: train on %" PRId64
241
+ " points after dedup "
242
+ "(was %" PRId64 " points)\n",
243
+ n2,
244
+ n);
252
245
  }
253
- IndexIVFFlat::train (n2, x2);
246
+ IndexIVFFlat::train(n2, x2.get());
254
247
  }
255
248
 
249
+ void IndexIVFFlatDedup::add_with_ids(
250
+ idx_t na,
251
+ const float* x,
252
+ const idx_t* xids) {
253
+ FAISS_THROW_IF_NOT(is_trained);
254
+ assert(invlists);
255
+ FAISS_THROW_IF_NOT_MSG(
256
+ direct_map.no(), "IVFFlatDedup not implemented with direct_map");
257
+ std::unique_ptr<int64_t[]> idx(new int64_t[na]);
258
+ quantizer->assign(na, x, idx.get());
256
259
 
260
+ int64_t n_add = 0, n_dup = 0;
257
261
 
258
- void IndexIVFFlatDedup::add_with_ids(
259
- idx_t na, const float* x, const idx_t* xids)
260
- {
262
+ #pragma omp parallel reduction(+ : n_add, n_dup)
263
+ {
264
+ int nt = omp_get_num_threads();
265
+ int rank = omp_get_thread_num();
261
266
 
262
- FAISS_THROW_IF_NOT (is_trained);
263
- assert (invlists);
264
- FAISS_THROW_IF_NOT_MSG (direct_map.no(),
265
- "IVFFlatDedup not implemented with direct_map");
266
- int64_t * idx = new int64_t [na];
267
- ScopeDeleter<int64_t> del (idx);
268
- quantizer->assign (na, x, idx);
267
+ // each thread takes care of a subset of lists
268
+ for (size_t i = 0; i < na; i++) {
269
+ int64_t list_no = idx[i];
269
270
 
270
- int64_t n_add = 0, n_dup = 0;
271
- // TODO make a omp loop with this
272
- for (size_t i = 0; i < na; i++) {
273
- idx_t id = xids ? xids[i] : ntotal + i;
274
- int64_t list_no = idx [i];
271
+ if (list_no < 0 || list_no % nt != rank) {
272
+ continue;
273
+ }
275
274
 
276
- if (list_no < 0) {
277
- continue;
278
- }
279
- const float *xi = x + i * d;
275
+ idx_t id = xids ? xids[i] : ntotal + i;
276
+ const float* xi = x + i * d;
280
277
 
281
- // search if there is already an entry with that id
282
- InvertedLists::ScopedCodes codes (invlists, list_no);
278
+ // search if there is already an entry with that id
279
+ InvertedLists::ScopedCodes codes(invlists, list_no);
283
280
 
284
- int64_t n = invlists->list_size (list_no);
285
- int64_t offset = -1;
286
- for (int64_t o = 0; o < n; o++) {
287
- if (!memcmp (codes.get() + o * code_size,
288
- xi, code_size)) {
289
- offset = o;
290
- break;
281
+ int64_t n = invlists->list_size(list_no);
282
+ int64_t offset = -1;
283
+ for (int64_t o = 0; o < n; o++) {
284
+ if (!memcmp(codes.get() + o * code_size, xi, code_size)) {
285
+ offset = o;
286
+ break;
287
+ }
291
288
  }
292
- }
293
289
 
294
- if (offset == -1) { // not found
295
- invlists->add_entry (list_no, id, (const uint8_t*) xi);
296
- } else {
297
- // mark equivalence
298
- idx_t id2 = invlists->get_single_id (list_no, offset);
299
- std::pair<idx_t, idx_t> pair (id2, id);
300
- instances.insert (pair);
301
- n_dup ++;
290
+ if (offset == -1) { // not found
291
+ invlists->add_entry(list_no, id, (const uint8_t*)xi);
292
+ } else {
293
+ // mark equivalence
294
+ idx_t id2 = invlists->get_single_id(list_no, offset);
295
+ std::pair<idx_t, idx_t> pair(id2, id);
296
+
297
+ #pragma omp critical
298
+ // executed by one thread at a time
299
+ instances.insert(pair);
300
+
301
+ n_dup++;
302
+ }
303
+ n_add++;
302
304
  }
303
- n_add++;
304
305
  }
305
306
  if (verbose) {
306
- printf("IndexIVFFlat::add_with_ids: added %" PRId64 " / %" PRId64 " vectors"
307
+ printf("IndexIVFFlat::add_with_ids: added %" PRId64 " / %" PRId64
308
+ " vectors"
307
309
  " (out of which %" PRId64 " are duplicates)\n",
308
- n_add, na, n_dup);
310
+ n_add,
311
+ na,
312
+ n_dup);
309
313
  }
310
314
  ntotal += n_add;
311
315
  }
312
316
 
313
- void IndexIVFFlatDedup::search_preassigned (
314
- idx_t n, const float *x, idx_t k,
315
- const idx_t *assign,
316
- const float *centroid_dis,
317
- float *distances, idx_t *labels,
318
- bool store_pairs,
319
- const IVFSearchParameters *params,
320
- IndexIVFStats *stats) const
321
- {
322
- FAISS_THROW_IF_NOT_MSG (
323
- !store_pairs, "store_pairs not supported in IVFDedup");
324
-
325
- IndexIVFFlat::search_preassigned (n, x, k, assign, centroid_dis,
326
- distances, labels, false,
327
- params);
328
-
329
- std::vector <idx_t> labels2 (k);
330
- std::vector <float> dis2 (k);
317
+ void IndexIVFFlatDedup::search_preassigned(
318
+ idx_t n,
319
+ const float* x,
320
+ idx_t k,
321
+ const idx_t* assign,
322
+ const float* centroid_dis,
323
+ float* distances,
324
+ idx_t* labels,
325
+ bool store_pairs,
326
+ const IVFSearchParameters* params,
327
+ IndexIVFStats* stats) const {
328
+ FAISS_THROW_IF_NOT_MSG(
329
+ !store_pairs, "store_pairs not supported in IVFDedup");
330
+
331
+ IndexIVFFlat::search_preassigned(
332
+ n, x, k, assign, centroid_dis, distances, labels, false, params);
333
+
334
+ std::vector<idx_t> labels2(k);
335
+ std::vector<float> dis2(k);
331
336
 
332
337
  for (int64_t i = 0; i < n; i++) {
333
- idx_t *labels1 = labels + i * k;
334
- float *dis1 = distances + i * k;
338
+ idx_t* labels1 = labels + i * k;
339
+ float* dis1 = distances + i * k;
335
340
  int64_t j = 0;
336
341
  for (; j < k; j++) {
337
- if (instances.find (labels1[j]) != instances.end ()) {
342
+ if (instances.find(labels1[j]) != instances.end()) {
338
343
  // a duplicate: special handling
339
344
  break;
340
345
  }
@@ -344,11 +349,11 @@ void IndexIVFFlatDedup::search_preassigned (
344
349
  int64_t j0 = j;
345
350
  int64_t rp = j;
346
351
  while (j < k) {
347
- auto range = instances.equal_range (labels1[rp]);
352
+ auto range = instances.equal_range(labels1[rp]);
348
353
  float dis = dis1[rp];
349
354
  labels2[j] = labels1[rp];
350
355
  dis2[j] = dis;
351
- j ++;
356
+ j++;
352
357
  for (auto it = range.first; j < k && it != range.second; ++it) {
353
358
  labels2[j] = it->second;
354
359
  dis2[j] = dis;
@@ -356,21 +361,18 @@ void IndexIVFFlatDedup::search_preassigned (
356
361
  }
357
362
  rp++;
358
363
  }
359
- memcpy (labels1 + j0, labels2.data() + j0,
360
- sizeof(labels1[0]) * (k - j0));
361
- memcpy (dis1 + j0, dis2.data() + j0,
362
- sizeof(dis2[0]) * (k - j0));
364
+ memcpy(labels1 + j0,
365
+ labels2.data() + j0,
366
+ sizeof(labels1[0]) * (k - j0));
367
+ memcpy(dis1 + j0, dis2.data() + j0, sizeof(dis2[0]) * (k - j0));
363
368
  }
364
369
  }
365
-
366
370
  }
367
371
 
368
-
369
- size_t IndexIVFFlatDedup::remove_ids(const IDSelector& sel)
370
- {
372
+ size_t IndexIVFFlatDedup::remove_ids(const IDSelector& sel) {
371
373
  std::unordered_map<idx_t, idx_t> replace;
372
- std::vector<std::pair<idx_t, idx_t> > toadd;
373
- for (auto it = instances.begin(); it != instances.end(); ) {
374
+ std::vector<std::pair<idx_t, idx_t>> toadd;
375
+ for (auto it = instances.begin(); it != instances.end();) {
374
376
  if (sel.is_member(it->first)) {
375
377
  // then we erase this entry
376
378
  if (!sel.is_member(it->second)) {
@@ -378,8 +380,8 @@ size_t IndexIVFFlatDedup::remove_ids(const IDSelector& sel)
378
380
  if (replace.count(it->first) == 0) {
379
381
  replace[it->first] = it->second;
380
382
  } else { // remember we should add an element
381
- std::pair<idx_t, idx_t> new_entry (
382
- replace[it->first], it->second);
383
+ std::pair<idx_t, idx_t> new_entry(
384
+ replace[it->first], it->second);
383
385
  toadd.push_back(new_entry);
384
386
  }
385
387
  }
@@ -393,32 +395,34 @@ size_t IndexIVFFlatDedup::remove_ids(const IDSelector& sel)
393
395
  }
394
396
  }
395
397
 
396
- instances.insert (toadd.begin(), toadd.end());
398
+ instances.insert(toadd.begin(), toadd.end());
397
399
 
398
400
  // mostly copied from IndexIVF.cpp
399
401
 
400
- FAISS_THROW_IF_NOT_MSG (direct_map.no(),
401
- "direct map remove not implemented");
402
+ FAISS_THROW_IF_NOT_MSG(
403
+ direct_map.no(), "direct map remove not implemented");
402
404
 
403
405
  std::vector<int64_t> toremove(nlist);
404
406
 
405
407
  #pragma omp parallel for
406
408
  for (int64_t i = 0; i < nlist; i++) {
407
- int64_t l0 = invlists->list_size (i), l = l0, j = 0;
408
- InvertedLists::ScopedIds idsi (invlists, i);
409
+ int64_t l0 = invlists->list_size(i), l = l0, j = 0;
410
+ InvertedLists::ScopedIds idsi(invlists, i);
409
411
  while (j < l) {
410
- if (sel.is_member (idsi[j])) {
412
+ if (sel.is_member(idsi[j])) {
411
413
  if (replace.count(idsi[j]) == 0) {
412
414
  l--;
413
- invlists->update_entry (
414
- i, j,
415
- invlists->get_single_id (i, l),
416
- InvertedLists::ScopedCodes (invlists, i, l).get());
415
+ invlists->update_entry(
416
+ i,
417
+ j,
418
+ invlists->get_single_id(i, l),
419
+ InvertedLists::ScopedCodes(invlists, i, l).get());
417
420
  } else {
418
- invlists->update_entry (
419
- i, j,
420
- replace[idsi[j]],
421
- InvertedLists::ScopedCodes (invlists, i, j).get());
421
+ invlists->update_entry(
422
+ i,
423
+ j,
424
+ replace[idsi[j]],
425
+ InvertedLists::ScopedCodes(invlists, i, j).get());
422
426
  j++;
423
427
  }
424
428
  } else {
@@ -432,37 +436,28 @@ size_t IndexIVFFlatDedup::remove_ids(const IDSelector& sel)
432
436
  for (int64_t i = 0; i < nlist; i++) {
433
437
  if (toremove[i] > 0) {
434
438
  nremove += toremove[i];
435
- invlists->resize(
436
- i, invlists->list_size(i) - toremove[i]);
439
+ invlists->resize(i, invlists->list_size(i) - toremove[i]);
437
440
  }
438
441
  }
439
442
  ntotal -= nremove;
440
443
  return nremove;
441
444
  }
442
445
 
443
-
444
446
  void IndexIVFFlatDedup::range_search(
445
- idx_t ,
446
- const float* ,
447
- float ,
448
- RangeSearchResult* ) const
449
- {
450
- FAISS_THROW_MSG ("not implemented");
447
+ idx_t,
448
+ const float*,
449
+ float,
450
+ RangeSearchResult*) const {
451
+ FAISS_THROW_MSG("not implemented");
451
452
  }
452
453
 
453
- void IndexIVFFlatDedup::update_vectors (int , const idx_t *, const float *)
454
- {
455
- FAISS_THROW_MSG ("not implemented");
454
+ void IndexIVFFlatDedup::update_vectors(int, const idx_t*, const float*) {
455
+ FAISS_THROW_MSG("not implemented");
456
456
  }
457
457
 
458
-
459
- void IndexIVFFlatDedup::reconstruct_from_offset (
460
- int64_t , int64_t , float* ) const
461
- {
462
- FAISS_THROW_MSG ("not implemented");
458
+ void IndexIVFFlatDedup::reconstruct_from_offset(int64_t, int64_t, float*)
459
+ const {
460
+ FAISS_THROW_MSG("not implemented");
463
461
  }
464
462
 
465
-
466
-
467
-
468
463
  } // namespace faiss