faiss 0.2.0 → 0.2.4

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 (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
@@ -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