faiss 0.1.7 → 0.2.3

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (219) hide show
  1. checksums.yaml +4 -4
  2. data/CHANGELOG.md +18 -0
  3. data/README.md +7 -7
  4. data/ext/faiss/ext.cpp +1 -1
  5. data/ext/faiss/extconf.rb +8 -2
  6. data/ext/faiss/index.cpp +102 -69
  7. data/ext/faiss/index_binary.cpp +24 -30
  8. data/ext/faiss/kmeans.cpp +20 -16
  9. data/ext/faiss/numo.hpp +867 -0
  10. data/ext/faiss/pca_matrix.cpp +13 -14
  11. data/ext/faiss/product_quantizer.cpp +23 -24
  12. data/ext/faiss/utils.cpp +10 -37
  13. data/ext/faiss/utils.h +2 -13
  14. data/lib/faiss/version.rb +1 -1
  15. data/lib/faiss.rb +0 -5
  16. data/vendor/faiss/faiss/AutoTune.cpp +292 -291
  17. data/vendor/faiss/faiss/AutoTune.h +55 -56
  18. data/vendor/faiss/faiss/Clustering.cpp +334 -195
  19. data/vendor/faiss/faiss/Clustering.h +88 -35
  20. data/vendor/faiss/faiss/IVFlib.cpp +171 -195
  21. data/vendor/faiss/faiss/IVFlib.h +48 -51
  22. data/vendor/faiss/faiss/Index.cpp +85 -103
  23. data/vendor/faiss/faiss/Index.h +54 -48
  24. data/vendor/faiss/faiss/Index2Layer.cpp +139 -164
  25. data/vendor/faiss/faiss/Index2Layer.h +22 -22
  26. data/vendor/faiss/faiss/IndexBinary.cpp +45 -37
  27. data/vendor/faiss/faiss/IndexBinary.h +140 -132
  28. data/vendor/faiss/faiss/IndexBinaryFlat.cpp +73 -53
  29. data/vendor/faiss/faiss/IndexBinaryFlat.h +29 -24
  30. data/vendor/faiss/faiss/IndexBinaryFromFloat.cpp +46 -43
  31. data/vendor/faiss/faiss/IndexBinaryFromFloat.h +16 -15
  32. data/vendor/faiss/faiss/IndexBinaryHNSW.cpp +215 -232
  33. data/vendor/faiss/faiss/IndexBinaryHNSW.h +25 -24
  34. data/vendor/faiss/faiss/IndexBinaryHash.cpp +182 -177
  35. data/vendor/faiss/faiss/IndexBinaryHash.h +41 -34
  36. data/vendor/faiss/faiss/IndexBinaryIVF.cpp +489 -461
  37. data/vendor/faiss/faiss/IndexBinaryIVF.h +97 -68
  38. data/vendor/faiss/faiss/IndexFlat.cpp +116 -147
  39. data/vendor/faiss/faiss/IndexFlat.h +35 -46
  40. data/vendor/faiss/faiss/IndexHNSW.cpp +372 -348
  41. data/vendor/faiss/faiss/IndexHNSW.h +57 -41
  42. data/vendor/faiss/faiss/IndexIVF.cpp +474 -454
  43. data/vendor/faiss/faiss/IndexIVF.h +146 -113
  44. data/vendor/faiss/faiss/IndexIVFFlat.cpp +248 -250
  45. data/vendor/faiss/faiss/IndexIVFFlat.h +48 -51
  46. data/vendor/faiss/faiss/IndexIVFPQ.cpp +457 -516
  47. data/vendor/faiss/faiss/IndexIVFPQ.h +74 -66
  48. data/vendor/faiss/faiss/IndexIVFPQFastScan.cpp +406 -372
  49. data/vendor/faiss/faiss/IndexIVFPQFastScan.h +82 -57
  50. data/vendor/faiss/faiss/IndexIVFPQR.cpp +104 -102
  51. data/vendor/faiss/faiss/IndexIVFPQR.h +33 -28
  52. data/vendor/faiss/faiss/IndexIVFSpectralHash.cpp +125 -133
  53. data/vendor/faiss/faiss/IndexIVFSpectralHash.h +19 -21
  54. data/vendor/faiss/faiss/IndexLSH.cpp +75 -96
  55. data/vendor/faiss/faiss/IndexLSH.h +21 -26
  56. data/vendor/faiss/faiss/IndexLattice.cpp +42 -56
  57. data/vendor/faiss/faiss/IndexLattice.h +11 -16
  58. data/vendor/faiss/faiss/IndexNNDescent.cpp +231 -0
  59. data/vendor/faiss/faiss/IndexNNDescent.h +72 -0
  60. data/vendor/faiss/faiss/IndexNSG.cpp +303 -0
  61. data/vendor/faiss/faiss/IndexNSG.h +85 -0
  62. data/vendor/faiss/faiss/IndexPQ.cpp +405 -464
  63. data/vendor/faiss/faiss/IndexPQ.h +64 -67
  64. data/vendor/faiss/faiss/IndexPQFastScan.cpp +143 -170
  65. data/vendor/faiss/faiss/IndexPQFastScan.h +46 -32
  66. data/vendor/faiss/faiss/IndexPreTransform.cpp +120 -150
  67. data/vendor/faiss/faiss/IndexPreTransform.h +33 -36
  68. data/vendor/faiss/faiss/IndexRefine.cpp +115 -131
  69. data/vendor/faiss/faiss/IndexRefine.h +22 -23
  70. data/vendor/faiss/faiss/IndexReplicas.cpp +147 -153
  71. data/vendor/faiss/faiss/IndexReplicas.h +62 -56
  72. data/vendor/faiss/faiss/IndexResidual.cpp +291 -0
  73. data/vendor/faiss/faiss/IndexResidual.h +152 -0
  74. data/vendor/faiss/faiss/IndexScalarQuantizer.cpp +120 -155
  75. data/vendor/faiss/faiss/IndexScalarQuantizer.h +41 -45
  76. data/vendor/faiss/faiss/IndexShards.cpp +256 -240
  77. data/vendor/faiss/faiss/IndexShards.h +85 -73
  78. data/vendor/faiss/faiss/MatrixStats.cpp +112 -97
  79. data/vendor/faiss/faiss/MatrixStats.h +7 -10
  80. data/vendor/faiss/faiss/MetaIndexes.cpp +135 -157
  81. data/vendor/faiss/faiss/MetaIndexes.h +40 -34
  82. data/vendor/faiss/faiss/MetricType.h +7 -7
  83. data/vendor/faiss/faiss/VectorTransform.cpp +652 -474
  84. data/vendor/faiss/faiss/VectorTransform.h +61 -89
  85. data/vendor/faiss/faiss/clone_index.cpp +77 -73
  86. data/vendor/faiss/faiss/clone_index.h +4 -9
  87. data/vendor/faiss/faiss/gpu/GpuAutoTune.cpp +33 -38
  88. data/vendor/faiss/faiss/gpu/GpuAutoTune.h +11 -9
  89. data/vendor/faiss/faiss/gpu/GpuCloner.cpp +197 -170
  90. data/vendor/faiss/faiss/gpu/GpuCloner.h +53 -35
  91. data/vendor/faiss/faiss/gpu/GpuClonerOptions.cpp +12 -14
  92. data/vendor/faiss/faiss/gpu/GpuClonerOptions.h +27 -25
  93. data/vendor/faiss/faiss/gpu/GpuDistance.h +116 -112
  94. data/vendor/faiss/faiss/gpu/GpuFaissAssert.h +1 -2
  95. data/vendor/faiss/faiss/gpu/GpuIndex.h +134 -137
  96. data/vendor/faiss/faiss/gpu/GpuIndexBinaryFlat.h +76 -73
  97. data/vendor/faiss/faiss/gpu/GpuIndexFlat.h +173 -162
  98. data/vendor/faiss/faiss/gpu/GpuIndexIVF.h +67 -64
  99. data/vendor/faiss/faiss/gpu/GpuIndexIVFFlat.h +89 -86
  100. data/vendor/faiss/faiss/gpu/GpuIndexIVFPQ.h +150 -141
  101. data/vendor/faiss/faiss/gpu/GpuIndexIVFScalarQuantizer.h +101 -103
  102. data/vendor/faiss/faiss/gpu/GpuIndicesOptions.h +17 -16
  103. data/vendor/faiss/faiss/gpu/GpuResources.cpp +116 -128
  104. data/vendor/faiss/faiss/gpu/GpuResources.h +182 -186
  105. data/vendor/faiss/faiss/gpu/StandardGpuResources.cpp +433 -422
  106. data/vendor/faiss/faiss/gpu/StandardGpuResources.h +131 -130
  107. data/vendor/faiss/faiss/gpu/impl/InterleavedCodes.cpp +468 -456
  108. data/vendor/faiss/faiss/gpu/impl/InterleavedCodes.h +25 -19
  109. data/vendor/faiss/faiss/gpu/impl/RemapIndices.cpp +22 -20
  110. data/vendor/faiss/faiss/gpu/impl/RemapIndices.h +9 -8
  111. data/vendor/faiss/faiss/gpu/perf/IndexWrapper-inl.h +39 -44
  112. data/vendor/faiss/faiss/gpu/perf/IndexWrapper.h +16 -14
  113. data/vendor/faiss/faiss/gpu/perf/PerfClustering.cpp +77 -71
  114. data/vendor/faiss/faiss/gpu/perf/PerfIVFPQAdd.cpp +109 -88
  115. data/vendor/faiss/faiss/gpu/perf/WriteIndex.cpp +75 -64
  116. data/vendor/faiss/faiss/gpu/test/TestCodePacking.cpp +230 -215
  117. data/vendor/faiss/faiss/gpu/test/TestGpuIndexBinaryFlat.cpp +80 -86
  118. data/vendor/faiss/faiss/gpu/test/TestGpuIndexFlat.cpp +284 -277
  119. data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFFlat.cpp +416 -416
  120. data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFPQ.cpp +611 -517
  121. data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFScalarQuantizer.cpp +166 -164
  122. data/vendor/faiss/faiss/gpu/test/TestGpuMemoryException.cpp +61 -53
  123. data/vendor/faiss/faiss/gpu/test/TestUtils.cpp +274 -238
  124. data/vendor/faiss/faiss/gpu/test/TestUtils.h +73 -57
  125. data/vendor/faiss/faiss/gpu/test/demo_ivfpq_indexing_gpu.cpp +47 -50
  126. data/vendor/faiss/faiss/gpu/utils/DeviceUtils.h +79 -72
  127. data/vendor/faiss/faiss/gpu/utils/StackDeviceMemory.cpp +140 -146
  128. data/vendor/faiss/faiss/gpu/utils/StackDeviceMemory.h +69 -71
  129. data/vendor/faiss/faiss/gpu/utils/StaticUtils.h +21 -16
  130. data/vendor/faiss/faiss/gpu/utils/Timer.cpp +25 -29
  131. data/vendor/faiss/faiss/gpu/utils/Timer.h +30 -29
  132. data/vendor/faiss/faiss/impl/AdditiveQuantizer.cpp +270 -0
  133. data/vendor/faiss/faiss/impl/AdditiveQuantizer.h +115 -0
  134. data/vendor/faiss/faiss/impl/AuxIndexStructures.cpp +90 -120
  135. data/vendor/faiss/faiss/impl/AuxIndexStructures.h +81 -65
  136. data/vendor/faiss/faiss/impl/FaissAssert.h +73 -58
  137. data/vendor/faiss/faiss/impl/FaissException.cpp +56 -48
  138. data/vendor/faiss/faiss/impl/FaissException.h +41 -29
  139. data/vendor/faiss/faiss/impl/HNSW.cpp +595 -611
  140. data/vendor/faiss/faiss/impl/HNSW.h +179 -200
  141. data/vendor/faiss/faiss/impl/LocalSearchQuantizer.cpp +672 -0
  142. data/vendor/faiss/faiss/impl/LocalSearchQuantizer.h +172 -0
  143. data/vendor/faiss/faiss/impl/NNDescent.cpp +487 -0
  144. data/vendor/faiss/faiss/impl/NNDescent.h +154 -0
  145. data/vendor/faiss/faiss/impl/NSG.cpp +682 -0
  146. data/vendor/faiss/faiss/impl/NSG.h +199 -0
  147. data/vendor/faiss/faiss/impl/PolysemousTraining.cpp +484 -454
  148. data/vendor/faiss/faiss/impl/PolysemousTraining.h +52 -55
  149. data/vendor/faiss/faiss/impl/ProductQuantizer-inl.h +26 -47
  150. data/vendor/faiss/faiss/impl/ProductQuantizer.cpp +469 -459
  151. data/vendor/faiss/faiss/impl/ProductQuantizer.h +76 -87
  152. data/vendor/faiss/faiss/impl/ResidualQuantizer.cpp +448 -0
  153. data/vendor/faiss/faiss/impl/ResidualQuantizer.h +130 -0
  154. data/vendor/faiss/faiss/impl/ResultHandler.h +96 -132
  155. data/vendor/faiss/faiss/impl/ScalarQuantizer.cpp +648 -701
  156. data/vendor/faiss/faiss/impl/ScalarQuantizer.h +48 -46
  157. data/vendor/faiss/faiss/impl/ThreadedIndex-inl.h +129 -131
  158. data/vendor/faiss/faiss/impl/ThreadedIndex.h +61 -55
  159. data/vendor/faiss/faiss/impl/index_read.cpp +547 -479
  160. data/vendor/faiss/faiss/impl/index_write.cpp +497 -407
  161. data/vendor/faiss/faiss/impl/io.cpp +75 -94
  162. data/vendor/faiss/faiss/impl/io.h +31 -41
  163. data/vendor/faiss/faiss/impl/io_macros.h +40 -29
  164. data/vendor/faiss/faiss/impl/lattice_Zn.cpp +137 -186
  165. data/vendor/faiss/faiss/impl/lattice_Zn.h +40 -51
  166. data/vendor/faiss/faiss/impl/platform_macros.h +29 -8
  167. data/vendor/faiss/faiss/impl/pq4_fast_scan.cpp +77 -124
  168. data/vendor/faiss/faiss/impl/pq4_fast_scan.h +39 -48
  169. data/vendor/faiss/faiss/impl/pq4_fast_scan_search_1.cpp +41 -52
  170. data/vendor/faiss/faiss/impl/pq4_fast_scan_search_qbs.cpp +80 -117
  171. data/vendor/faiss/faiss/impl/simd_result_handlers.h +109 -137
  172. data/vendor/faiss/faiss/index_factory.cpp +269 -218
  173. data/vendor/faiss/faiss/index_factory.h +6 -7
  174. data/vendor/faiss/faiss/index_io.h +23 -26
  175. data/vendor/faiss/faiss/invlists/BlockInvertedLists.cpp +67 -75
  176. data/vendor/faiss/faiss/invlists/BlockInvertedLists.h +22 -24
  177. data/vendor/faiss/faiss/invlists/DirectMap.cpp +96 -112
  178. data/vendor/faiss/faiss/invlists/DirectMap.h +29 -33
  179. data/vendor/faiss/faiss/invlists/InvertedLists.cpp +307 -364
  180. data/vendor/faiss/faiss/invlists/InvertedLists.h +151 -151
  181. data/vendor/faiss/faiss/invlists/InvertedListsIOHook.cpp +29 -34
  182. data/vendor/faiss/faiss/invlists/InvertedListsIOHook.h +17 -18
  183. data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.cpp +257 -293
  184. data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.h +50 -45
  185. data/vendor/faiss/faiss/python/python_callbacks.cpp +23 -26
  186. data/vendor/faiss/faiss/python/python_callbacks.h +9 -16
  187. data/vendor/faiss/faiss/utils/AlignedTable.h +79 -44
  188. data/vendor/faiss/faiss/utils/Heap.cpp +40 -48
  189. data/vendor/faiss/faiss/utils/Heap.h +186 -209
  190. data/vendor/faiss/faiss/utils/WorkerThread.cpp +67 -76
  191. data/vendor/faiss/faiss/utils/WorkerThread.h +32 -33
  192. data/vendor/faiss/faiss/utils/distances.cpp +301 -310
  193. data/vendor/faiss/faiss/utils/distances.h +133 -118
  194. data/vendor/faiss/faiss/utils/distances_simd.cpp +456 -516
  195. data/vendor/faiss/faiss/utils/extra_distances-inl.h +117 -0
  196. data/vendor/faiss/faiss/utils/extra_distances.cpp +113 -232
  197. data/vendor/faiss/faiss/utils/extra_distances.h +30 -29
  198. data/vendor/faiss/faiss/utils/hamming-inl.h +260 -209
  199. data/vendor/faiss/faiss/utils/hamming.cpp +375 -469
  200. data/vendor/faiss/faiss/utils/hamming.h +62 -85
  201. data/vendor/faiss/faiss/utils/ordered_key_value.h +16 -18
  202. data/vendor/faiss/faiss/utils/partitioning.cpp +393 -318
  203. data/vendor/faiss/faiss/utils/partitioning.h +26 -21
  204. data/vendor/faiss/faiss/utils/quantize_lut.cpp +78 -66
  205. data/vendor/faiss/faiss/utils/quantize_lut.h +22 -20
  206. data/vendor/faiss/faiss/utils/random.cpp +39 -63
  207. data/vendor/faiss/faiss/utils/random.h +13 -16
  208. data/vendor/faiss/faiss/utils/simdlib.h +4 -2
  209. data/vendor/faiss/faiss/utils/simdlib_avx2.h +88 -85
  210. data/vendor/faiss/faiss/utils/simdlib_emulated.h +226 -165
  211. data/vendor/faiss/faiss/utils/simdlib_neon.h +832 -0
  212. data/vendor/faiss/faiss/utils/utils.cpp +304 -287
  213. data/vendor/faiss/faiss/utils/utils.h +53 -48
  214. metadata +26 -12
  215. data/lib/faiss/index.rb +0 -20
  216. data/lib/faiss/index_binary.rb +0 -20
  217. data/lib/faiss/kmeans.rb +0 -15
  218. data/lib/faiss/pca_matrix.rb +0 -15
  219. data/lib/faiss/product_quantizer.rb +0 -22
@@ -15,7 +15,6 @@
15
15
 
16
16
  namespace faiss {
17
17
 
18
-
19
18
  /** Fast scan version of IVFPQ. Works for 4-bit PQ for now.
20
19
  *
21
20
  * The codes in the inverted lists are not stored sequentially but
@@ -32,13 +31,12 @@ namespace faiss {
32
31
  * 13: idem, collect results in reservoir
33
32
  */
34
33
 
35
- struct IndexIVFPQFastScan: IndexIVF {
36
-
37
- bool by_residual; ///< Encode residual or plain vector?
38
- ProductQuantizer pq; ///< produces the codes
34
+ struct IndexIVFPQFastScan : IndexIVF {
35
+ bool by_residual; ///< Encode residual or plain vector?
36
+ ProductQuantizer pq; ///< produces the codes
39
37
 
40
38
  // size of the kernel
41
- int bbs; // set at build time
39
+ int bbs; // set at build time
42
40
 
43
41
  // M rounded up to a multiple of 2
44
42
  size_t M2;
@@ -57,88 +55,114 @@ struct IndexIVFPQFastScan: IndexIVF {
57
55
  int qbs = 0;
58
56
  size_t qbs2 = 0;
59
57
 
60
- IndexIVFPQFastScan (
61
- Index * quantizer, size_t d, size_t nlist,
62
- size_t M, size_t nbits_per_idx,
63
- MetricType metric = METRIC_L2, int bbs = 32);
58
+ IndexIVFPQFastScan(
59
+ Index* quantizer,
60
+ size_t d,
61
+ size_t nlist,
62
+ size_t M,
63
+ size_t nbits_per_idx,
64
+ MetricType metric = METRIC_L2,
65
+ int bbs = 32);
64
66
 
65
- IndexIVFPQFastScan ();
67
+ IndexIVFPQFastScan();
66
68
 
67
69
  // built from an IndexIVFPQ
68
- explicit IndexIVFPQFastScan(const IndexIVFPQ & orig, int bbs = 32);
70
+ explicit IndexIVFPQFastScan(const IndexIVFPQ& orig, int bbs = 32);
69
71
 
70
72
  /// orig's inverted lists (for debugging)
71
- InvertedLists * orig_invlists = nullptr;
73
+ InvertedLists* orig_invlists = nullptr;
72
74
 
73
- void train_residual (idx_t n, const float *x) override;
75
+ void train_residual(idx_t n, const float* x) override;
74
76
 
75
77
  /// build precomputed table, possibly updating use_precomputed_table
76
- void precompute_table ();
78
+ void precompute_table();
77
79
 
78
80
  /// same as the regular IVFPQ encoder. The codes are not reorganized by
79
81
  /// blocks a that point
80
82
  void encode_vectors(
81
- idx_t n, const float* x,
82
- const idx_t *list_nos, uint8_t * codes,
83
+ idx_t n,
84
+ const float* x,
85
+ const idx_t* list_nos,
86
+ uint8_t* codes,
83
87
  bool include_listno = false) const override;
84
88
 
85
- void add_with_ids (
86
- idx_t n, const float * x, const idx_t *xids) override;
89
+ void add_with_ids(idx_t n, const float* x, const idx_t* xids) override;
87
90
 
88
91
  void search(
89
- idx_t n, const float* x, idx_t k,
90
- float* distances, idx_t* labels) const override;
92
+ idx_t n,
93
+ const float* x,
94
+ idx_t k,
95
+ float* distances,
96
+ idx_t* labels) const override;
91
97
 
92
98
  // prepare look-up tables
93
99
 
94
100
  void compute_LUT(
95
- size_t n, const float *x,
96
- const idx_t *coarse_ids, const float *coarse_dis,
97
- AlignedTable<float> & dis_tables,
98
- AlignedTable<float> & biases
99
- ) const;
101
+ size_t n,
102
+ const float* x,
103
+ const idx_t* coarse_ids,
104
+ const float* coarse_dis,
105
+ AlignedTable<float>& dis_tables,
106
+ AlignedTable<float>& biases) const;
100
107
 
101
108
  void compute_LUT_uint8(
102
- size_t n, const float *x,
103
- const idx_t *coarse_ids, const float *coarse_dis,
104
- AlignedTable<uint8_t> & dis_tables,
105
- AlignedTable<uint16_t> & biases,
106
- float * normalizers
107
- ) const;
109
+ size_t n,
110
+ const float* x,
111
+ const idx_t* coarse_ids,
112
+ const float* coarse_dis,
113
+ AlignedTable<uint8_t>& dis_tables,
114
+ AlignedTable<uint16_t>& biases,
115
+ float* normalizers) const;
108
116
 
109
117
  // internal search funcs
110
118
 
111
- template<bool is_max>
119
+ template <bool is_max>
112
120
  void search_dispatch_implem(
113
- idx_t n, const float* x, idx_t k,
114
- float* distances, idx_t* labels) const;
121
+ idx_t n,
122
+ const float* x,
123
+ idx_t k,
124
+ float* distances,
125
+ idx_t* labels) const;
115
126
 
116
- template<class C>
127
+ template <class C>
117
128
  void search_implem_1(
118
- idx_t n, const float* x, idx_t k,
119
- float* distances, idx_t* labels) const;
129
+ idx_t n,
130
+ const float* x,
131
+ idx_t k,
132
+ float* distances,
133
+ idx_t* labels) const;
120
134
 
121
- template<class C>
135
+ template <class C>
122
136
  void search_implem_2(
123
- idx_t n, const float* x, idx_t k,
124
- float* distances, idx_t* labels) const;
137
+ idx_t n,
138
+ const float* x,
139
+ idx_t k,
140
+ float* distances,
141
+ idx_t* labels) const;
125
142
 
126
143
  // implem 10 and 12 are not multithreaded internally, so
127
144
  // export search stats
128
- template<class C>
145
+ template <class C>
129
146
  void search_implem_10(
130
- idx_t n, const float* x, idx_t k,
131
- float* distances, idx_t* labels,
132
- int impl, size_t *ndis_out, size_t *nlist_out) const;
133
-
134
- template<class C>
147
+ idx_t n,
148
+ const float* x,
149
+ idx_t k,
150
+ float* distances,
151
+ idx_t* labels,
152
+ int impl,
153
+ size_t* ndis_out,
154
+ size_t* nlist_out) const;
155
+
156
+ template <class C>
135
157
  void search_implem_12(
136
- idx_t n, const float* x, idx_t k,
137
- float* distances, idx_t* labels,
138
- int impl, size_t *ndis_out, size_t *nlist_out) const;
139
-
140
-
141
-
158
+ idx_t n,
159
+ const float* x,
160
+ idx_t k,
161
+ float* distances,
162
+ idx_t* labels,
163
+ int impl,
164
+ size_t* ndis_out,
165
+ size_t* nlist_out) const;
142
166
  };
143
167
 
144
168
  struct IVFFastScanStats {
@@ -148,13 +172,15 @@ struct IVFFastScanStats {
148
172
  uint64_t reservoir_times[4];
149
173
 
150
174
  double Mcy_at(int i) {
151
- return times[i] / (1000*1000.0);
175
+ return times[i] / (1000 * 1000.0);
152
176
  }
153
177
 
154
178
  double Mcy_reservoir_at(int i) {
155
- return reservoir_times[i] / (1000*1000.0);
179
+ return reservoir_times[i] / (1000 * 1000.0);
180
+ }
181
+ IVFFastScanStats() {
182
+ reset();
156
183
  }
157
- IVFFastScanStats() {reset();}
158
184
  void reset() {
159
185
  memset(this, 0, sizeof(*this));
160
186
  }
@@ -162,5 +188,4 @@ struct IVFFastScanStats {
162
188
 
163
189
  FAISS_API extern IVFFastScanStats IVFFastScan_stats;
164
190
 
165
-
166
191
  } // namespace faiss
@@ -12,115 +12,116 @@
12
12
  #include <cinttypes>
13
13
 
14
14
  #include <faiss/utils/Heap.h>
15
- #include <faiss/utils/utils.h>
16
15
  #include <faiss/utils/distances.h>
16
+ #include <faiss/utils/utils.h>
17
17
 
18
18
  #include <faiss/impl/FaissAssert.h>
19
19
 
20
-
21
20
  namespace faiss {
22
21
 
23
22
  /*****************************************
24
23
  * IndexIVFPQR implementation
25
24
  ******************************************/
26
25
 
27
- IndexIVFPQR::IndexIVFPQR (
28
- Index * quantizer, size_t d, size_t nlist,
29
- size_t M, size_t nbits_per_idx,
30
- size_t M_refine, size_t nbits_per_idx_refine):
31
- IndexIVFPQ (quantizer, d, nlist, M, nbits_per_idx),
32
- refine_pq (d, M_refine, nbits_per_idx_refine),
33
- k_factor (4)
34
- {
26
+ IndexIVFPQR::IndexIVFPQR(
27
+ Index* quantizer,
28
+ size_t d,
29
+ size_t nlist,
30
+ size_t M,
31
+ size_t nbits_per_idx,
32
+ size_t M_refine,
33
+ size_t nbits_per_idx_refine)
34
+ : IndexIVFPQ(quantizer, d, nlist, M, nbits_per_idx),
35
+ refine_pq(d, M_refine, nbits_per_idx_refine),
36
+ k_factor(4) {
35
37
  by_residual = true;
36
38
  }
37
39
 
38
- IndexIVFPQR::IndexIVFPQR ():
39
- k_factor (1)
40
- {
40
+ IndexIVFPQR::IndexIVFPQR() : k_factor(1) {
41
41
  by_residual = true;
42
42
  }
43
43
 
44
-
45
-
46
- void IndexIVFPQR::reset()
47
- {
44
+ void IndexIVFPQR::reset() {
48
45
  IndexIVFPQ::reset();
49
46
  refine_codes.clear();
50
47
  }
51
48
 
49
+ void IndexIVFPQR::train_residual(idx_t n, const float* x) {
50
+ float* residual_2 = new float[n * d];
51
+ ScopeDeleter<float> del(residual_2);
52
52
 
53
-
54
-
55
- void IndexIVFPQR::train_residual (idx_t n, const float *x)
56
- {
57
-
58
- float * residual_2 = new float [n * d];
59
- ScopeDeleter <float> del(residual_2);
60
-
61
- train_residual_o (n, x, residual_2);
53
+ train_residual_o(n, x, residual_2);
62
54
 
63
55
  if (verbose)
64
- printf ("training %zdx%zd 2nd level PQ quantizer on %" PRId64 " %dD-vectors\n",
65
- refine_pq.M, refine_pq.ksub, n, d);
56
+ printf("training %zdx%zd 2nd level PQ quantizer on %" PRId64
57
+ " %dD-vectors\n",
58
+ refine_pq.M,
59
+ refine_pq.ksub,
60
+ n,
61
+ d);
66
62
 
67
63
  refine_pq.cp.max_points_per_centroid = 1000;
68
64
  refine_pq.cp.verbose = verbose;
69
65
 
70
- refine_pq.train (n, residual_2);
71
-
66
+ refine_pq.train(n, residual_2);
72
67
  }
73
68
 
74
-
75
- void IndexIVFPQR::add_with_ids (idx_t n, const float *x, const idx_t *xids) {
76
- add_core (n, x, xids, nullptr);
69
+ void IndexIVFPQR::add_with_ids(idx_t n, const float* x, const idx_t* xids) {
70
+ add_core(n, x, xids, nullptr);
77
71
  }
78
72
 
79
- void IndexIVFPQR::add_core (idx_t n, const float *x, const idx_t *xids,
80
- const idx_t *precomputed_idx) {
81
-
82
- float * residual_2 = new float [n * d];
83
- ScopeDeleter <float> del(residual_2);
73
+ void IndexIVFPQR::add_core(
74
+ idx_t n,
75
+ const float* x,
76
+ const idx_t* xids,
77
+ const idx_t* precomputed_idx) {
78
+ float* residual_2 = new float[n * d];
79
+ ScopeDeleter<float> del(residual_2);
84
80
 
85
81
  idx_t n0 = ntotal;
86
82
 
87
- add_core_o (n, x, xids, residual_2, precomputed_idx);
88
-
89
- refine_codes.resize (ntotal * refine_pq.code_size);
90
-
91
- refine_pq.compute_codes (
92
- residual_2, &refine_codes[n0 * refine_pq.code_size], n);
83
+ add_core_o(n, x, xids, residual_2, precomputed_idx);
93
84
 
85
+ refine_codes.resize(ntotal * refine_pq.code_size);
94
86
 
87
+ refine_pq.compute_codes(
88
+ residual_2, &refine_codes[n0 * refine_pq.code_size], n);
95
89
  }
96
90
  #define TIC t0 = get_cycles()
97
- #define TOC get_cycles () - t0
98
-
99
-
100
- void IndexIVFPQR::search_preassigned (
101
- idx_t n, const float *x, idx_t k,
102
- const idx_t *idx, const float *L1_dis,
103
- float *distances, idx_t *labels,
91
+ #define TOC get_cycles() - t0
92
+
93
+ void IndexIVFPQR::search_preassigned(
94
+ idx_t n,
95
+ const float* x,
96
+ idx_t k,
97
+ const idx_t* idx,
98
+ const float* L1_dis,
99
+ float* distances,
100
+ idx_t* labels,
104
101
  bool store_pairs,
105
- const IVFSearchParameters *params, IndexIVFStats *stats
106
- ) const
107
- {
102
+ const IVFSearchParameters* params,
103
+ IndexIVFStats* stats) const {
108
104
  uint64_t t0;
109
105
  TIC;
110
106
  size_t k_coarse = long(k * k_factor);
111
- idx_t *coarse_labels = new idx_t [k_coarse * n];
112
- ScopeDeleter<idx_t> del1 (coarse_labels);
107
+ idx_t* coarse_labels = new idx_t[k_coarse * n];
108
+ ScopeDeleter<idx_t> del1(coarse_labels);
113
109
  { // query with quantizer levels 1 and 2.
114
- float *coarse_distances = new float [k_coarse * n];
110
+ float* coarse_distances = new float[k_coarse * n];
115
111
  ScopeDeleter<float> del(coarse_distances);
116
112
 
117
- IndexIVFPQ::search_preassigned (
118
- n, x, k_coarse,
119
- idx, L1_dis, coarse_distances, coarse_labels,
120
- true, params);
113
+ IndexIVFPQ::search_preassigned(
114
+ n,
115
+ x,
116
+ k_coarse,
117
+ idx,
118
+ L1_dis,
119
+ coarse_distances,
120
+ coarse_labels,
121
+ true,
122
+ params);
121
123
  }
122
124
 
123
-
124
125
  indexIVFPQ_stats.search_cycles += TOC;
125
126
 
126
127
  TIC;
@@ -130,91 +131,92 @@ void IndexIVFPQR::search_preassigned (
130
131
  #pragma omp parallel reduction(+ : n_refine)
131
132
  {
132
133
  // tmp buffers
133
- float *residual_1 = new float [2 * d];
134
- ScopeDeleter<float> del (residual_1);
135
- float *residual_2 = residual_1 + d;
134
+ float* residual_1 = new float[2 * d];
135
+ ScopeDeleter<float> del(residual_1);
136
+ float* residual_2 = residual_1 + d;
136
137
  #pragma omp for
137
138
  for (idx_t i = 0; i < n; i++) {
138
- const float *xq = x + i * d;
139
- const idx_t * shortlist = coarse_labels + k_coarse * i;
140
- float * heap_sim = distances + k * i;
141
- idx_t * heap_ids = labels + k * i;
142
- maxheap_heapify (k, heap_sim, heap_ids);
139
+ const float* xq = x + i * d;
140
+ const idx_t* shortlist = coarse_labels + k_coarse * i;
141
+ float* heap_sim = distances + k * i;
142
+ idx_t* heap_ids = labels + k * i;
143
+ maxheap_heapify(k, heap_sim, heap_ids);
143
144
 
144
145
  for (int j = 0; j < k_coarse; j++) {
145
146
  idx_t sl = shortlist[j];
146
147
 
147
- if (sl == -1) continue;
148
+ if (sl == -1)
149
+ continue;
148
150
 
149
151
  int list_no = lo_listno(sl);
150
152
  int ofs = lo_offset(sl);
151
153
 
152
- assert (list_no >= 0 && list_no < nlist);
153
- assert (ofs >= 0 && ofs < invlists->list_size (list_no));
154
+ assert(list_no >= 0 && list_no < nlist);
155
+ assert(ofs >= 0 && ofs < invlists->list_size(list_no));
154
156
 
155
157
  // 1st level residual
156
- quantizer->compute_residual (xq, residual_1, list_no);
158
+ quantizer->compute_residual(xq, residual_1, list_no);
157
159
 
158
160
  // 2nd level residual
159
- const uint8_t * l2code =
160
- invlists->get_single_code (list_no, ofs);
161
+ const uint8_t* l2code = invlists->get_single_code(list_no, ofs);
161
162
 
162
- pq.decode (l2code, residual_2);
163
+ pq.decode(l2code, residual_2);
163
164
  for (int l = 0; l < d; l++)
164
165
  residual_2[l] = residual_1[l] - residual_2[l];
165
166
 
166
167
  // 3rd level residual's approximation
167
- idx_t id = invlists->get_single_id (list_no, ofs);
168
- assert (0 <= id && id < ntotal);
169
- refine_pq.decode (&refine_codes [id * refine_pq.code_size],
170
- residual_1);
168
+ idx_t id = invlists->get_single_id(list_no, ofs);
169
+ assert(0 <= id && id < ntotal);
170
+ refine_pq.decode(
171
+ &refine_codes[id * refine_pq.code_size], residual_1);
171
172
 
172
- float dis = fvec_L2sqr (residual_1, residual_2, d);
173
+ float dis = fvec_L2sqr(residual_1, residual_2, d);
173
174
 
174
175
  if (dis < heap_sim[0]) {
175
176
  idx_t id_or_pair = store_pairs ? sl : id;
176
- maxheap_replace_top (k, heap_sim, heap_ids, dis, id_or_pair);
177
+ maxheap_replace_top(k, heap_sim, heap_ids, dis, id_or_pair);
177
178
  }
178
- n_refine ++;
179
+ n_refine++;
179
180
  }
180
- maxheap_reorder (k, heap_sim, heap_ids);
181
+ maxheap_reorder(k, heap_sim, heap_ids);
181
182
  }
182
183
  }
183
184
  indexIVFPQ_stats.nrefine += n_refine;
184
185
  indexIVFPQ_stats.refine_cycles += TOC;
185
186
  }
186
187
 
187
- void IndexIVFPQR::reconstruct_from_offset (int64_t list_no, int64_t offset,
188
- float* recons) const
189
- {
190
- IndexIVFPQ::reconstruct_from_offset (list_no, offset, recons);
188
+ void IndexIVFPQR::reconstruct_from_offset(
189
+ int64_t list_no,
190
+ int64_t offset,
191
+ float* recons) const {
192
+ IndexIVFPQ::reconstruct_from_offset(list_no, offset, recons);
191
193
 
192
- idx_t id = invlists->get_single_id (list_no, offset);
193
- assert (0 <= id && id < ntotal);
194
+ idx_t id = invlists->get_single_id(list_no, offset);
195
+ assert(0 <= id && id < ntotal);
194
196
 
195
197
  std::vector<float> r3(d);
196
- refine_pq.decode (&refine_codes [id * refine_pq.code_size], r3.data());
198
+ refine_pq.decode(&refine_codes[id * refine_pq.code_size], r3.data());
197
199
  for (int i = 0; i < d; ++i) {
198
- recons[i] += r3[i];
200
+ recons[i] += r3[i];
199
201
  }
200
202
  }
201
203
 
202
- void IndexIVFPQR::merge_from (IndexIVF &other_in, idx_t add_id)
203
- {
204
- IndexIVFPQR *other = dynamic_cast<IndexIVFPQR *> (&other_in);
204
+ void IndexIVFPQR::merge_from(IndexIVF& other_in, idx_t add_id) {
205
+ IndexIVFPQR* other = dynamic_cast<IndexIVFPQR*>(&other_in);
205
206
  FAISS_THROW_IF_NOT(other);
206
207
 
207
- IndexIVF::merge_from (other_in, add_id);
208
+ IndexIVF::merge_from(other_in, add_id);
208
209
 
209
- refine_codes.insert (refine_codes.end(),
210
- other->refine_codes.begin(),
211
- other->refine_codes.end());
210
+ refine_codes.insert(
211
+ refine_codes.end(),
212
+ other->refine_codes.begin(),
213
+ other->refine_codes.end());
212
214
  other->refine_codes.clear();
213
215
  }
214
216
 
215
217
  size_t IndexIVFPQR::remove_ids(const IDSelector& /*sel*/) {
216
- FAISS_THROW_MSG("not implemented");
217
- return 0;
218
+ FAISS_THROW_MSG("not implemented");
219
+ return 0;
218
220
  }
219
221
 
220
222
  } // namespace faiss
@@ -13,23 +13,24 @@
13
13
 
14
14
  #include <faiss/IndexIVFPQ.h>
15
15
 
16
-
17
16
  namespace faiss {
18
17
 
19
-
20
-
21
18
  /** Index with an additional level of PQ refinement */
22
- struct IndexIVFPQR: IndexIVFPQ {
23
- ProductQuantizer refine_pq; ///< 3rd level quantizer
24
- std::vector <uint8_t> refine_codes; ///< corresponding codes
19
+ struct IndexIVFPQR : IndexIVFPQ {
20
+ ProductQuantizer refine_pq; ///< 3rd level quantizer
21
+ std::vector<uint8_t> refine_codes; ///< corresponding codes
25
22
 
26
23
  /// factor between k requested in search and the k requested from the IVFPQ
27
24
  float k_factor;
28
25
 
29
- IndexIVFPQR (
30
- Index * quantizer, size_t d, size_t nlist,
31
- size_t M, size_t nbits_per_idx,
32
- size_t M_refine, size_t nbits_per_idx_refine);
26
+ IndexIVFPQR(
27
+ Index* quantizer,
28
+ size_t d,
29
+ size_t nlist,
30
+ size_t M,
31
+ size_t nbits_per_idx,
32
+ size_t M_refine,
33
+ size_t nbits_per_idx_refine);
33
34
 
34
35
  void reset() override;
35
36
 
@@ -41,26 +42,30 @@ struct IndexIVFPQR: IndexIVFPQ {
41
42
  void add_with_ids(idx_t n, const float* x, const idx_t* xids) override;
42
43
 
43
44
  /// same as add_with_ids, but optionally use the precomputed list ids
44
- void add_core (idx_t n, const float *x, const idx_t *xids,
45
- const idx_t *precomputed_idx = nullptr);
46
-
47
- void reconstruct_from_offset (int64_t list_no, int64_t offset,
48
- float* recons) const override;
49
-
50
- void merge_from (IndexIVF &other, idx_t add_id) override;
51
-
52
-
53
- void search_preassigned (idx_t n, const float *x, idx_t k,
54
- const idx_t *assign,
55
- const float *centroid_dis,
56
- float *distances, idx_t *labels,
57
- bool store_pairs,
58
- const IVFSearchParameters *params=nullptr,
59
- IndexIVFStats *stats=nullptr
60
- ) const override;
45
+ void add_core(
46
+ idx_t n,
47
+ const float* x,
48
+ const idx_t* xids,
49
+ const idx_t* precomputed_idx) override;
50
+
51
+ void reconstruct_from_offset(int64_t list_no, int64_t offset, float* recons)
52
+ const override;
53
+
54
+ void merge_from(IndexIVF& other, idx_t add_id) override;
55
+
56
+ void search_preassigned(
57
+ idx_t n,
58
+ const float* x,
59
+ idx_t k,
60
+ const idx_t* assign,
61
+ const float* centroid_dis,
62
+ float* distances,
63
+ idx_t* labels,
64
+ bool store_pairs,
65
+ const IVFSearchParameters* params = nullptr,
66
+ IndexIVFStats* stats = nullptr) const override;
61
67
 
62
68
  IndexIVFPQR();
63
69
  };
64
70
 
65
-
66
71
  } // namespace faiss