faiss 0.3.0 → 0.3.2

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 (216) hide show
  1. checksums.yaml +4 -4
  2. data/CHANGELOG.md +9 -0
  3. data/LICENSE.txt +1 -1
  4. data/README.md +1 -1
  5. data/ext/faiss/extconf.rb +9 -2
  6. data/ext/faiss/index.cpp +1 -1
  7. data/ext/faiss/index_binary.cpp +2 -2
  8. data/ext/faiss/product_quantizer.cpp +1 -1
  9. data/lib/faiss/version.rb +1 -1
  10. data/vendor/faiss/faiss/AutoTune.cpp +7 -7
  11. data/vendor/faiss/faiss/AutoTune.h +1 -2
  12. data/vendor/faiss/faiss/Clustering.cpp +39 -22
  13. data/vendor/faiss/faiss/Clustering.h +40 -21
  14. data/vendor/faiss/faiss/IVFlib.cpp +26 -12
  15. data/vendor/faiss/faiss/Index.cpp +1 -1
  16. data/vendor/faiss/faiss/Index.h +40 -10
  17. data/vendor/faiss/faiss/Index2Layer.cpp +7 -7
  18. data/vendor/faiss/faiss/IndexAdditiveQuantizer.cpp +176 -166
  19. data/vendor/faiss/faiss/IndexAdditiveQuantizerFastScan.cpp +15 -15
  20. data/vendor/faiss/faiss/IndexBinary.cpp +9 -4
  21. data/vendor/faiss/faiss/IndexBinary.h +8 -19
  22. data/vendor/faiss/faiss/IndexBinaryFromFloat.cpp +2 -1
  23. data/vendor/faiss/faiss/IndexBinaryHNSW.cpp +24 -31
  24. data/vendor/faiss/faiss/IndexBinaryHNSW.h +1 -1
  25. data/vendor/faiss/faiss/IndexBinaryHash.cpp +25 -50
  26. data/vendor/faiss/faiss/IndexBinaryIVF.cpp +107 -188
  27. data/vendor/faiss/faiss/IndexFastScan.cpp +95 -146
  28. data/vendor/faiss/faiss/IndexFastScan.h +9 -8
  29. data/vendor/faiss/faiss/IndexFlat.cpp +206 -10
  30. data/vendor/faiss/faiss/IndexFlat.h +20 -1
  31. data/vendor/faiss/faiss/IndexFlatCodes.cpp +170 -5
  32. data/vendor/faiss/faiss/IndexFlatCodes.h +23 -4
  33. data/vendor/faiss/faiss/IndexHNSW.cpp +231 -382
  34. data/vendor/faiss/faiss/IndexHNSW.h +62 -49
  35. data/vendor/faiss/faiss/IndexIDMap.cpp +69 -28
  36. data/vendor/faiss/faiss/IndexIDMap.h +24 -2
  37. data/vendor/faiss/faiss/IndexIVF.cpp +162 -56
  38. data/vendor/faiss/faiss/IndexIVF.h +46 -6
  39. data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.cpp +33 -26
  40. data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.h +6 -2
  41. data/vendor/faiss/faiss/IndexIVFAdditiveQuantizerFastScan.cpp +19 -46
  42. data/vendor/faiss/faiss/IndexIVFAdditiveQuantizerFastScan.h +4 -3
  43. data/vendor/faiss/faiss/IndexIVFFastScan.cpp +502 -401
  44. data/vendor/faiss/faiss/IndexIVFFastScan.h +63 -26
  45. data/vendor/faiss/faiss/IndexIVFFlat.cpp +15 -5
  46. data/vendor/faiss/faiss/IndexIVFFlat.h +3 -2
  47. data/vendor/faiss/faiss/IndexIVFIndependentQuantizer.cpp +172 -0
  48. data/vendor/faiss/faiss/IndexIVFIndependentQuantizer.h +56 -0
  49. data/vendor/faiss/faiss/IndexIVFPQ.cpp +79 -125
  50. data/vendor/faiss/faiss/IndexIVFPQ.h +6 -7
  51. data/vendor/faiss/faiss/IndexIVFPQFastScan.cpp +39 -52
  52. data/vendor/faiss/faiss/IndexIVFPQFastScan.h +4 -3
  53. data/vendor/faiss/faiss/IndexIVFPQR.cpp +45 -29
  54. data/vendor/faiss/faiss/IndexIVFPQR.h +5 -2
  55. data/vendor/faiss/faiss/IndexIVFSpectralHash.cpp +25 -27
  56. data/vendor/faiss/faiss/IndexIVFSpectralHash.h +6 -6
  57. data/vendor/faiss/faiss/IndexLSH.cpp +14 -16
  58. data/vendor/faiss/faiss/IndexLattice.cpp +1 -19
  59. data/vendor/faiss/faiss/IndexLattice.h +3 -22
  60. data/vendor/faiss/faiss/IndexNNDescent.cpp +3 -33
  61. data/vendor/faiss/faiss/IndexNNDescent.h +1 -1
  62. data/vendor/faiss/faiss/IndexNSG.cpp +11 -27
  63. data/vendor/faiss/faiss/IndexNSG.h +11 -11
  64. data/vendor/faiss/faiss/IndexNeuralNetCodec.cpp +56 -0
  65. data/vendor/faiss/faiss/IndexNeuralNetCodec.h +49 -0
  66. data/vendor/faiss/faiss/IndexPQ.cpp +72 -88
  67. data/vendor/faiss/faiss/IndexPQ.h +1 -4
  68. data/vendor/faiss/faiss/IndexPQFastScan.cpp +1 -1
  69. data/vendor/faiss/faiss/IndexPreTransform.cpp +25 -31
  70. data/vendor/faiss/faiss/IndexPreTransform.h +1 -1
  71. data/vendor/faiss/faiss/IndexRefine.cpp +54 -24
  72. data/vendor/faiss/faiss/IndexRefine.h +7 -0
  73. data/vendor/faiss/faiss/IndexReplicas.cpp +23 -26
  74. data/vendor/faiss/faiss/IndexScalarQuantizer.cpp +25 -17
  75. data/vendor/faiss/faiss/IndexScalarQuantizer.h +6 -4
  76. data/vendor/faiss/faiss/IndexShards.cpp +21 -29
  77. data/vendor/faiss/faiss/IndexShardsIVF.cpp +1 -2
  78. data/vendor/faiss/faiss/MatrixStats.cpp +17 -32
  79. data/vendor/faiss/faiss/MatrixStats.h +21 -9
  80. data/vendor/faiss/faiss/MetaIndexes.cpp +35 -35
  81. data/vendor/faiss/faiss/MetricType.h +7 -2
  82. data/vendor/faiss/faiss/VectorTransform.cpp +13 -26
  83. data/vendor/faiss/faiss/VectorTransform.h +7 -7
  84. data/vendor/faiss/faiss/clone_index.cpp +15 -10
  85. data/vendor/faiss/faiss/clone_index.h +3 -0
  86. data/vendor/faiss/faiss/cppcontrib/detail/UintReader.h +95 -17
  87. data/vendor/faiss/faiss/cppcontrib/factory_tools.cpp +152 -0
  88. data/vendor/faiss/faiss/cppcontrib/factory_tools.h +24 -0
  89. data/vendor/faiss/faiss/cppcontrib/sa_decode/Level2-inl.h +83 -30
  90. data/vendor/faiss/faiss/gpu/GpuCloner.cpp +123 -8
  91. data/vendor/faiss/faiss/gpu/GpuCloner.h +22 -0
  92. data/vendor/faiss/faiss/gpu/GpuClonerOptions.h +13 -0
  93. data/vendor/faiss/faiss/gpu/GpuDistance.h +46 -38
  94. data/vendor/faiss/faiss/gpu/GpuFaissAssert.h +1 -1
  95. data/vendor/faiss/faiss/gpu/GpuIndex.h +30 -12
  96. data/vendor/faiss/faiss/gpu/GpuIndexCagra.h +282 -0
  97. data/vendor/faiss/faiss/gpu/GpuIndexFlat.h +4 -4
  98. data/vendor/faiss/faiss/gpu/GpuIndexIVF.h +14 -9
  99. data/vendor/faiss/faiss/gpu/GpuIndexIVFFlat.h +20 -3
  100. data/vendor/faiss/faiss/gpu/GpuIndexIVFPQ.h +22 -11
  101. data/vendor/faiss/faiss/gpu/GpuIndexIVFScalarQuantizer.h +1 -3
  102. data/vendor/faiss/faiss/gpu/GpuResources.cpp +24 -3
  103. data/vendor/faiss/faiss/gpu/GpuResources.h +39 -11
  104. data/vendor/faiss/faiss/gpu/StandardGpuResources.cpp +142 -17
  105. data/vendor/faiss/faiss/gpu/StandardGpuResources.h +57 -3
  106. data/vendor/faiss/faiss/gpu/impl/InterleavedCodes.cpp +26 -21
  107. data/vendor/faiss/faiss/gpu/perf/PerfClustering.cpp +7 -1
  108. data/vendor/faiss/faiss/gpu/test/TestCodePacking.cpp +8 -5
  109. data/vendor/faiss/faiss/gpu/test/TestGpuIndexBinaryFlat.cpp +25 -0
  110. data/vendor/faiss/faiss/gpu/test/TestGpuIndexFlat.cpp +129 -9
  111. data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFFlat.cpp +332 -40
  112. data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFPQ.cpp +299 -208
  113. data/vendor/faiss/faiss/gpu/test/TestGpuMemoryException.cpp +1 -0
  114. data/vendor/faiss/faiss/gpu/test/demo_ivfpq_indexing_gpu.cpp +1 -1
  115. data/vendor/faiss/faiss/gpu/utils/DeviceUtils.h +6 -0
  116. data/vendor/faiss/faiss/gpu/utils/RaftUtils.h +75 -0
  117. data/vendor/faiss/faiss/gpu/utils/Timer.cpp +4 -1
  118. data/vendor/faiss/faiss/gpu/utils/Timer.h +1 -1
  119. data/vendor/faiss/faiss/impl/AdditiveQuantizer.cpp +3 -1
  120. data/vendor/faiss/faiss/impl/AdditiveQuantizer.h +5 -5
  121. data/vendor/faiss/faiss/impl/AuxIndexStructures.cpp +26 -1
  122. data/vendor/faiss/faiss/impl/AuxIndexStructures.h +10 -3
  123. data/vendor/faiss/faiss/impl/DistanceComputer.h +70 -1
  124. data/vendor/faiss/faiss/impl/FaissAssert.h +4 -2
  125. data/vendor/faiss/faiss/impl/FaissException.h +13 -34
  126. data/vendor/faiss/faiss/impl/HNSW.cpp +605 -186
  127. data/vendor/faiss/faiss/impl/HNSW.h +52 -30
  128. data/vendor/faiss/faiss/impl/IDSelector.h +4 -4
  129. data/vendor/faiss/faiss/impl/LocalSearchQuantizer.cpp +11 -9
  130. data/vendor/faiss/faiss/impl/LookupTableScaler.h +34 -0
  131. data/vendor/faiss/faiss/impl/NNDescent.cpp +42 -27
  132. data/vendor/faiss/faiss/impl/NSG.cpp +0 -29
  133. data/vendor/faiss/faiss/impl/NSG.h +1 -1
  134. data/vendor/faiss/faiss/impl/PolysemousTraining.cpp +14 -12
  135. data/vendor/faiss/faiss/impl/ProductAdditiveQuantizer.h +1 -1
  136. data/vendor/faiss/faiss/impl/ProductQuantizer.cpp +25 -22
  137. data/vendor/faiss/faiss/impl/ProductQuantizer.h +6 -2
  138. data/vendor/faiss/faiss/impl/Quantizer.h +1 -1
  139. data/vendor/faiss/faiss/impl/ResidualQuantizer.cpp +27 -1015
  140. data/vendor/faiss/faiss/impl/ResidualQuantizer.h +5 -63
  141. data/vendor/faiss/faiss/impl/ResultHandler.h +347 -172
  142. data/vendor/faiss/faiss/impl/ScalarQuantizer.cpp +1104 -147
  143. data/vendor/faiss/faiss/impl/ScalarQuantizer.h +3 -8
  144. data/vendor/faiss/faiss/impl/code_distance/code_distance-avx2.h +285 -42
  145. data/vendor/faiss/faiss/impl/code_distance/code_distance-avx512.h +248 -0
  146. data/vendor/faiss/faiss/impl/code_distance/code_distance-generic.h +21 -14
  147. data/vendor/faiss/faiss/impl/code_distance/code_distance.h +22 -12
  148. data/vendor/faiss/faiss/impl/index_read.cpp +74 -34
  149. data/vendor/faiss/faiss/impl/index_read_utils.h +37 -0
  150. data/vendor/faiss/faiss/impl/index_write.cpp +88 -51
  151. data/vendor/faiss/faiss/impl/io.cpp +23 -15
  152. data/vendor/faiss/faiss/impl/io.h +4 -4
  153. data/vendor/faiss/faiss/impl/io_macros.h +6 -0
  154. data/vendor/faiss/faiss/impl/lattice_Zn.cpp +1 -1
  155. data/vendor/faiss/faiss/impl/platform_macros.h +40 -1
  156. data/vendor/faiss/faiss/impl/pq4_fast_scan.cpp +14 -0
  157. data/vendor/faiss/faiss/impl/pq4_fast_scan.h +7 -6
  158. data/vendor/faiss/faiss/impl/pq4_fast_scan_search_1.cpp +52 -38
  159. data/vendor/faiss/faiss/impl/pq4_fast_scan_search_qbs.cpp +487 -49
  160. data/vendor/faiss/faiss/impl/residual_quantizer_encode_steps.cpp +960 -0
  161. data/vendor/faiss/faiss/impl/residual_quantizer_encode_steps.h +176 -0
  162. data/vendor/faiss/faiss/impl/simd_result_handlers.h +481 -225
  163. data/vendor/faiss/faiss/index_factory.cpp +41 -20
  164. data/vendor/faiss/faiss/index_io.h +12 -5
  165. data/vendor/faiss/faiss/invlists/BlockInvertedLists.cpp +28 -8
  166. data/vendor/faiss/faiss/invlists/BlockInvertedLists.h +3 -0
  167. data/vendor/faiss/faiss/invlists/DirectMap.cpp +10 -2
  168. data/vendor/faiss/faiss/invlists/InvertedLists.cpp +73 -17
  169. data/vendor/faiss/faiss/invlists/InvertedLists.h +26 -8
  170. data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.cpp +24 -9
  171. data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.h +2 -1
  172. data/vendor/faiss/faiss/python/python_callbacks.cpp +4 -4
  173. data/vendor/faiss/faiss/utils/Heap.cpp +3 -1
  174. data/vendor/faiss/faiss/utils/Heap.h +105 -0
  175. data/vendor/faiss/faiss/utils/NeuralNet.cpp +342 -0
  176. data/vendor/faiss/faiss/utils/NeuralNet.h +147 -0
  177. data/vendor/faiss/faiss/utils/WorkerThread.h +1 -0
  178. data/vendor/faiss/faiss/utils/bf16.h +36 -0
  179. data/vendor/faiss/faiss/utils/distances.cpp +147 -123
  180. data/vendor/faiss/faiss/utils/distances.h +86 -9
  181. data/vendor/faiss/faiss/utils/distances_fused/avx512.cpp +5 -5
  182. data/vendor/faiss/faiss/utils/distances_fused/avx512.h +2 -2
  183. data/vendor/faiss/faiss/utils/distances_fused/distances_fused.cpp +2 -2
  184. data/vendor/faiss/faiss/utils/distances_fused/distances_fused.h +1 -1
  185. data/vendor/faiss/faiss/utils/distances_fused/simdlib_based.cpp +5 -5
  186. data/vendor/faiss/faiss/utils/distances_fused/simdlib_based.h +1 -1
  187. data/vendor/faiss/faiss/utils/distances_simd.cpp +1589 -243
  188. data/vendor/faiss/faiss/utils/extra_distances-inl.h +70 -0
  189. data/vendor/faiss/faiss/utils/extra_distances.cpp +85 -137
  190. data/vendor/faiss/faiss/utils/extra_distances.h +3 -2
  191. data/vendor/faiss/faiss/utils/fp16-arm.h +29 -0
  192. data/vendor/faiss/faiss/utils/fp16.h +2 -0
  193. data/vendor/faiss/faiss/utils/hamming.cpp +163 -111
  194. data/vendor/faiss/faiss/utils/hamming.h +58 -0
  195. data/vendor/faiss/faiss/utils/hamming_distance/avx2-inl.h +16 -89
  196. data/vendor/faiss/faiss/utils/hamming_distance/common.h +1 -0
  197. data/vendor/faiss/faiss/utils/hamming_distance/generic-inl.h +19 -88
  198. data/vendor/faiss/faiss/utils/hamming_distance/hamdis-inl.h +58 -0
  199. data/vendor/faiss/faiss/utils/hamming_distance/neon-inl.h +14 -104
  200. data/vendor/faiss/faiss/utils/partitioning.cpp +3 -4
  201. data/vendor/faiss/faiss/utils/prefetch.h +77 -0
  202. data/vendor/faiss/faiss/utils/quantize_lut.cpp +0 -14
  203. data/vendor/faiss/faiss/utils/random.cpp +43 -0
  204. data/vendor/faiss/faiss/utils/random.h +25 -0
  205. data/vendor/faiss/faiss/utils/simdlib.h +10 -1
  206. data/vendor/faiss/faiss/utils/simdlib_avx2.h +0 -6
  207. data/vendor/faiss/faiss/utils/simdlib_avx512.h +296 -0
  208. data/vendor/faiss/faiss/utils/simdlib_neon.h +77 -79
  209. data/vendor/faiss/faiss/utils/simdlib_ppc64.h +1084 -0
  210. data/vendor/faiss/faiss/utils/sorting.cpp +140 -5
  211. data/vendor/faiss/faiss/utils/sorting.h +27 -0
  212. data/vendor/faiss/faiss/utils/transpose/transpose-avx512-inl.h +176 -0
  213. data/vendor/faiss/faiss/utils/utils.cpp +120 -7
  214. data/vendor/faiss/faiss/utils/utils.h +60 -20
  215. metadata +23 -4
  216. data/vendor/faiss/faiss/impl/code_distance/code_distance_avx512.h +0 -102
@@ -4,10 +4,26 @@
4
4
  * This source code is licensed under the MIT license found in the
5
5
  * LICENSE file in the root directory of this source tree.
6
6
  */
7
+ /*
8
+ * Copyright (c) 2023, NVIDIA CORPORATION.
9
+ *
10
+ * Licensed under the Apache License, Version 2.0 (the "License");
11
+ * you may not use this file except in compliance with the License.
12
+ * You may obtain a copy of the License at
13
+ *
14
+ * http://www.apache.org/licenses/LICENSE-2.0
15
+ *
16
+ * Unless required by applicable law or agreed to in writing, software
17
+ * distributed under the License is distributed on an "AS IS" BASIS,
18
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
19
+ * See the License for the specific language governing permissions and
20
+ * limitations under the License.
21
+ */
7
22
 
8
23
  #include <faiss/IndexFlat.h>
9
24
  #include <faiss/IndexIVFFlat.h>
10
25
  #include <faiss/gpu/GpuIndexIVFFlat.h>
26
+ #include <faiss/gpu/GpuIndicesOptions.h>
11
27
  #include <faiss/gpu/StandardGpuResources.h>
12
28
  #include <faiss/gpu/test/TestUtils.h>
13
29
  #include <faiss/gpu/utils/DeviceUtils.h>
@@ -40,6 +56,8 @@ struct Options {
40
56
  faiss::gpu::INDICES_64_BIT});
41
57
 
42
58
  device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
59
+
60
+ useRaft = false;
43
61
  }
44
62
 
45
63
  std::string toString() const {
@@ -47,7 +65,7 @@ struct Options {
47
65
  str << "IVFFlat device " << device << " numVecs " << numAdd << " dim "
48
66
  << dim << " numCentroids " << numCentroids << " nprobe " << nprobe
49
67
  << " numQuery " << numQuery << " k " << k << " indicesOpt "
50
- << indicesOpt;
68
+ << indicesOpt << " useRaft " << useRaft;
51
69
 
52
70
  return str.str();
53
71
  }
@@ -61,6 +79,7 @@ struct Options {
61
79
  int k;
62
80
  int device;
63
81
  faiss::gpu::IndicesOptions indicesOpt;
82
+ bool useRaft;
64
83
  };
65
84
 
66
85
  void queryTest(
@@ -91,6 +110,7 @@ void queryTest(
91
110
  config.device = opt.device;
92
111
  config.indicesOptions = opt.indicesOpt;
93
112
  config.flatConfig.useFloat16 = useFloat16CoarseQuantizer;
113
+ config.use_raft = opt.useRaft;
94
114
 
95
115
  faiss::gpu::GpuIndexIVFFlat gpuIndex(
96
116
  &res, cpuIndex.d, cpuIndex.nlist, cpuIndex.metric_type, config);
@@ -114,7 +134,10 @@ void queryTest(
114
134
  }
115
135
  }
116
136
 
117
- void addTest(faiss::MetricType metricType, bool useFloat16CoarseQuantizer) {
137
+ void addTest(
138
+ faiss::MetricType metricType,
139
+ bool useFloat16CoarseQuantizer,
140
+ bool useRaft) {
118
141
  for (int tries = 0; tries < 2; ++tries) {
119
142
  Options opt;
120
143
 
@@ -138,8 +161,10 @@ void addTest(faiss::MetricType metricType, bool useFloat16CoarseQuantizer) {
138
161
 
139
162
  faiss::gpu::GpuIndexIVFFlatConfig config;
140
163
  config.device = opt.device;
141
- config.indicesOptions = opt.indicesOpt;
164
+ config.indicesOptions =
165
+ useRaft ? faiss::gpu::INDICES_64_BIT : opt.indicesOpt;
142
166
  config.flatConfig.useFloat16 = useFloat16CoarseQuantizer;
167
+ config.use_raft = useRaft;
143
168
 
144
169
  faiss::gpu::GpuIndexIVFFlat gpuIndex(
145
170
  &res, cpuIndex.d, cpuIndex.nlist, cpuIndex.metric_type, config);
@@ -163,7 +188,7 @@ void addTest(faiss::MetricType metricType, bool useFloat16CoarseQuantizer) {
163
188
  }
164
189
  }
165
190
 
166
- void copyToTest(bool useFloat16CoarseQuantizer) {
191
+ void copyToTest(bool useFloat16CoarseQuantizer, bool useRaft) {
167
192
  Options opt;
168
193
  std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
169
194
  std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
@@ -173,8 +198,10 @@ void copyToTest(bool useFloat16CoarseQuantizer) {
173
198
 
174
199
  faiss::gpu::GpuIndexIVFFlatConfig config;
175
200
  config.device = opt.device;
176
- config.indicesOptions = opt.indicesOpt;
201
+ config.indicesOptions =
202
+ useRaft ? faiss::gpu::INDICES_64_BIT : opt.indicesOpt;
177
203
  config.flatConfig.useFloat16 = useFloat16CoarseQuantizer;
204
+ config.use_raft = useRaft;
178
205
 
179
206
  faiss::gpu::GpuIndexIVFFlat gpuIndex(
180
207
  &res, opt.dim, opt.numCentroids, faiss::METRIC_L2, config);
@@ -214,7 +241,7 @@ void copyToTest(bool useFloat16CoarseQuantizer) {
214
241
  compFloat16 ? 0.30f : 0.015f);
215
242
  }
216
243
 
217
- void copyFromTest(bool useFloat16CoarseQuantizer) {
244
+ void copyFromTest(bool useFloat16CoarseQuantizer, bool useRaft) {
218
245
  Options opt;
219
246
  std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
220
247
  std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
@@ -232,8 +259,10 @@ void copyFromTest(bool useFloat16CoarseQuantizer) {
232
259
 
233
260
  faiss::gpu::GpuIndexIVFFlatConfig config;
234
261
  config.device = opt.device;
235
- config.indicesOptions = opt.indicesOpt;
262
+ config.indicesOptions =
263
+ useRaft ? faiss::gpu::INDICES_64_BIT : opt.indicesOpt;
236
264
  config.flatConfig.useFloat16 = useFloat16CoarseQuantizer;
265
+ config.use_raft = useRaft;
237
266
 
238
267
  faiss::gpu::GpuIndexIVFFlat gpuIndex(&res, 1, 1, faiss::METRIC_L2, config);
239
268
  gpuIndex.nprobe = 1;
@@ -265,19 +294,35 @@ void copyFromTest(bool useFloat16CoarseQuantizer) {
265
294
  }
266
295
 
267
296
  TEST(TestGpuIndexIVFFlat, Float32_32_Add_L2) {
268
- addTest(faiss::METRIC_L2, false);
297
+ addTest(faiss::METRIC_L2, false, false);
298
+
299
+ #if defined USE_NVIDIA_RAFT
300
+ addTest(faiss::METRIC_L2, false, true);
301
+ #endif
269
302
  }
270
303
 
271
304
  TEST(TestGpuIndexIVFFlat, Float32_32_Add_IP) {
272
- addTest(faiss::METRIC_INNER_PRODUCT, false);
305
+ addTest(faiss::METRIC_INNER_PRODUCT, false, false);
306
+
307
+ #if defined USE_NVIDIA_RAFT
308
+ addTest(faiss::METRIC_INNER_PRODUCT, false, true);
309
+ #endif
273
310
  }
274
311
 
275
312
  TEST(TestGpuIndexIVFFlat, Float16_32_Add_L2) {
276
- addTest(faiss::METRIC_L2, true);
313
+ addTest(faiss::METRIC_L2, true, false);
314
+
315
+ #if defined USE_NVIDIA_RAFT
316
+ addTest(faiss::METRIC_L2, true, true);
317
+ #endif
277
318
  }
278
319
 
279
320
  TEST(TestGpuIndexIVFFlat, Float16_32_Add_IP) {
280
- addTest(faiss::METRIC_INNER_PRODUCT, true);
321
+ addTest(faiss::METRIC_INNER_PRODUCT, true, false);
322
+
323
+ #if defined USE_NVIDIA_RAFT
324
+ addTest(faiss::METRIC_INNER_PRODUCT, true, true);
325
+ #endif
281
326
  }
282
327
 
283
328
  //
@@ -285,11 +330,25 @@ TEST(TestGpuIndexIVFFlat, Float16_32_Add_IP) {
285
330
  //
286
331
 
287
332
  TEST(TestGpuIndexIVFFlat, Float32_Query_L2) {
288
- queryTest(Options(), faiss::METRIC_L2, false);
333
+ Options opt;
334
+ queryTest(opt, faiss::METRIC_L2, false);
335
+
336
+ #if defined USE_NVIDIA_RAFT
337
+ opt.useRaft = true;
338
+ opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
339
+ queryTest(opt, faiss::METRIC_L2, false);
340
+ #endif
289
341
  }
290
342
 
291
343
  TEST(TestGpuIndexIVFFlat, Float32_Query_IP) {
292
- queryTest(Options(), faiss::METRIC_INNER_PRODUCT, false);
344
+ Options opt;
345
+ queryTest(opt, faiss::METRIC_INNER_PRODUCT, false);
346
+
347
+ #if defined USE_NVIDIA_RAFT
348
+ opt.useRaft = true;
349
+ opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
350
+ queryTest(opt, faiss::METRIC_INNER_PRODUCT, false);
351
+ #endif
293
352
  }
294
353
 
295
354
  TEST(TestGpuIndexIVFFlat, LargeBatch) {
@@ -297,16 +356,36 @@ TEST(TestGpuIndexIVFFlat, LargeBatch) {
297
356
  opt.dim = 3;
298
357
  opt.numQuery = 100000;
299
358
  queryTest(opt, faiss::METRIC_L2, false);
359
+
360
+ #if defined USE_NVIDIA_RAFT
361
+ opt.useRaft = true;
362
+ opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
363
+ queryTest(opt, faiss::METRIC_L2, false);
364
+ #endif
300
365
  }
301
366
 
302
367
  // float16 coarse quantizer
303
368
 
304
369
  TEST(TestGpuIndexIVFFlat, Float16_32_Query_L2) {
305
- queryTest(Options(), faiss::METRIC_L2, true);
370
+ Options opt;
371
+ queryTest(opt, faiss::METRIC_L2, true);
372
+
373
+ #if defined USE_NVIDIA_RAFT
374
+ opt.useRaft = true;
375
+ opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
376
+ queryTest(opt, faiss::METRIC_L2, true);
377
+ #endif
306
378
  }
307
379
 
308
380
  TEST(TestGpuIndexIVFFlat, Float16_32_Query_IP) {
309
- queryTest(Options(), faiss::METRIC_INNER_PRODUCT, true);
381
+ Options opt;
382
+ queryTest(opt, faiss::METRIC_INNER_PRODUCT, true);
383
+
384
+ #if defined USE_NVIDIA_RAFT
385
+ opt.useRaft = true;
386
+ opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
387
+ queryTest(opt, faiss::METRIC_INNER_PRODUCT, true);
388
+ #endif
310
389
  }
311
390
 
312
391
  //
@@ -318,24 +397,48 @@ TEST(TestGpuIndexIVFFlat, Float32_Query_L2_64) {
318
397
  Options opt;
319
398
  opt.dim = 64;
320
399
  queryTest(opt, faiss::METRIC_L2, false);
400
+
401
+ #if defined USE_NVIDIA_RAFT
402
+ opt.useRaft = true;
403
+ opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
404
+ queryTest(opt, faiss::METRIC_L2, false);
405
+ #endif
321
406
  }
322
407
 
323
408
  TEST(TestGpuIndexIVFFlat, Float32_Query_IP_64) {
324
409
  Options opt;
325
410
  opt.dim = 64;
326
411
  queryTest(opt, faiss::METRIC_INNER_PRODUCT, false);
412
+
413
+ #if defined USE_NVIDIA_RAFT
414
+ opt.useRaft = true;
415
+ opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
416
+ queryTest(opt, faiss::METRIC_INNER_PRODUCT, false);
417
+ #endif
327
418
  }
328
419
 
329
420
  TEST(TestGpuIndexIVFFlat, Float32_Query_L2_128) {
330
421
  Options opt;
331
422
  opt.dim = 128;
332
423
  queryTest(opt, faiss::METRIC_L2, false);
424
+
425
+ #if defined USE_NVIDIA_RAFT
426
+ opt.useRaft = true;
427
+ opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
428
+ queryTest(opt, faiss::METRIC_L2, false);
429
+ #endif
333
430
  }
334
431
 
335
432
  TEST(TestGpuIndexIVFFlat, Float32_Query_IP_128) {
336
433
  Options opt;
337
434
  opt.dim = 128;
338
435
  queryTest(opt, faiss::METRIC_INNER_PRODUCT, false);
436
+
437
+ #if defined USE_NVIDIA_RAFT
438
+ opt.useRaft = true;
439
+ opt.indicesOpt = faiss::gpu::INDICES_64_BIT;
440
+ queryTest(opt, faiss::METRIC_INNER_PRODUCT, false);
441
+ #endif
339
442
  }
340
443
 
341
444
  //
@@ -343,11 +446,19 @@ TEST(TestGpuIndexIVFFlat, Float32_Query_IP_128) {
343
446
  //
344
447
 
345
448
  TEST(TestGpuIndexIVFFlat, Float32_32_CopyTo) {
346
- copyToTest(false);
449
+ copyToTest(false, false);
450
+
451
+ #if defined USE_NVIDIA_RAFT
452
+ copyToTest(false, true);
453
+ #endif
347
454
  }
348
455
 
349
456
  TEST(TestGpuIndexIVFFlat, Float32_32_CopyFrom) {
350
- copyFromTest(false);
457
+ copyFromTest(false, false);
458
+
459
+ #if defined USE_NVIDIA_RAFT
460
+ copyFromTest(false, true);
461
+ #endif
351
462
  }
352
463
 
353
464
  TEST(TestGpuIndexIVFFlat, Float32_negative) {
@@ -377,23 +488,24 @@ TEST(TestGpuIndexIVFFlat, Float32_negative) {
377
488
  faiss::gpu::StandardGpuResources res;
378
489
  res.noTempMemory();
379
490
 
491
+ // Construct a positive test set
492
+ auto queryVecs = faiss::gpu::randVecs(opt.numQuery, opt.dim);
493
+
494
+ // Put all vecs on positive size
495
+ for (auto& f : queryVecs) {
496
+ f = std::abs(f);
497
+ }
498
+
380
499
  faiss::gpu::GpuIndexIVFFlatConfig config;
381
500
  config.device = opt.device;
382
501
  config.indicesOptions = opt.indicesOpt;
502
+ config.use_raft = false;
383
503
 
384
504
  faiss::gpu::GpuIndexIVFFlat gpuIndex(
385
505
  &res, cpuIndex.d, cpuIndex.nlist, cpuIndex.metric_type, config);
386
506
  gpuIndex.copyFrom(&cpuIndex);
387
507
  gpuIndex.nprobe = opt.nprobe;
388
508
 
389
- // Construct a positive test set
390
- auto queryVecs = faiss::gpu::randVecs(opt.numQuery, opt.dim);
391
-
392
- // Put all vecs on positive size
393
- for (auto& f : queryVecs) {
394
- f = std::abs(f);
395
- }
396
-
397
509
  bool compFloat16 = false;
398
510
  faiss::gpu::compareIndices(
399
511
  queryVecs,
@@ -409,6 +521,31 @@ TEST(TestGpuIndexIVFFlat, Float32_negative) {
409
521
  // in fp16. Figure out another way to test
410
522
  compFloat16 ? 0.99f : 0.1f,
411
523
  compFloat16 ? 0.65f : 0.015f);
524
+
525
+ #if defined USE_NVIDIA_RAFT
526
+ config.use_raft = true;
527
+ config.indicesOptions = faiss::gpu::INDICES_64_BIT;
528
+
529
+ faiss::gpu::GpuIndexIVFFlat raftGpuIndex(
530
+ &res, cpuIndex.d, cpuIndex.nlist, cpuIndex.metric_type, config);
531
+ raftGpuIndex.copyFrom(&cpuIndex);
532
+ raftGpuIndex.nprobe = opt.nprobe;
533
+
534
+ faiss::gpu::compareIndices(
535
+ queryVecs,
536
+ cpuIndex,
537
+ raftGpuIndex,
538
+ opt.numQuery,
539
+ opt.dim,
540
+ opt.k,
541
+ opt.toString(),
542
+ compFloat16 ? kF16MaxRelErr : kF32MaxRelErr,
543
+ // FIXME: the fp16 bounds are
544
+ // useless when math (the accumulator) is
545
+ // in fp16. Figure out another way to test
546
+ compFloat16 ? 0.99f : 0.1f,
547
+ compFloat16 ? 0.65f : 0.015f);
548
+ #endif
412
549
  }
413
550
 
414
551
  //
@@ -424,10 +561,18 @@ TEST(TestGpuIndexIVFFlat, QueryNaN) {
424
561
  faiss::gpu::StandardGpuResources res;
425
562
  res.noTempMemory();
426
563
 
564
+ int numQuery = 10;
565
+ std::vector<float> nans(
566
+ numQuery * opt.dim, std::numeric_limits<float>::quiet_NaN());
567
+
568
+ std::vector<float> distances(numQuery * opt.k, 0);
569
+ std::vector<faiss::idx_t> indices(numQuery * opt.k, 0);
570
+
427
571
  faiss::gpu::GpuIndexIVFFlatConfig config;
428
572
  config.device = opt.device;
429
573
  config.indicesOptions = opt.indicesOpt;
430
574
  config.flatConfig.useFloat16 = faiss::gpu::randBool();
575
+ config.use_raft = false;
431
576
 
432
577
  faiss::gpu::GpuIndexIVFFlat gpuIndex(
433
578
  &res, opt.dim, opt.numCentroids, faiss::METRIC_L2, config);
@@ -436,14 +581,31 @@ TEST(TestGpuIndexIVFFlat, QueryNaN) {
436
581
  gpuIndex.train(opt.numTrain, trainVecs.data());
437
582
  gpuIndex.add(opt.numAdd, addVecs.data());
438
583
 
439
- int numQuery = 10;
440
- std::vector<float> nans(
441
- numQuery * opt.dim, std::numeric_limits<float>::quiet_NaN());
584
+ gpuIndex.search(
585
+ numQuery, nans.data(), opt.k, distances.data(), indices.data());
442
586
 
443
- std::vector<float> distances(numQuery * opt.k, 0);
444
- std::vector<faiss::idx_t> indices(numQuery * opt.k, 0);
587
+ for (int q = 0; q < numQuery; ++q) {
588
+ for (int k = 0; k < opt.k; ++k) {
589
+ EXPECT_EQ(indices[q * opt.k + k], -1);
590
+ EXPECT_EQ(
591
+ distances[q * opt.k + k],
592
+ std::numeric_limits<float>::max());
593
+ }
594
+ }
445
595
 
446
- gpuIndex.search(
596
+ #if defined USE_NVIDIA_RAFT
597
+ config.use_raft = true;
598
+ config.indicesOptions = faiss::gpu::INDICES_64_BIT;
599
+ std::fill(distances.begin(), distances.end(), 0);
600
+ std::fill(indices.begin(), indices.end(), 0);
601
+ faiss::gpu::GpuIndexIVFFlat raftGpuIndex(
602
+ &res, opt.dim, opt.numCentroids, faiss::METRIC_L2, config);
603
+ raftGpuIndex.nprobe = opt.nprobe;
604
+
605
+ raftGpuIndex.train(opt.numTrain, trainVecs.data());
606
+ raftGpuIndex.add(opt.numAdd, addVecs.data());
607
+
608
+ raftGpuIndex.search(
447
609
  numQuery, nans.data(), opt.k, distances.data(), indices.data());
448
610
 
449
611
  for (int q = 0; q < numQuery; ++q) {
@@ -454,6 +616,7 @@ TEST(TestGpuIndexIVFFlat, QueryNaN) {
454
616
  std::numeric_limits<float>::max());
455
617
  }
456
618
  }
619
+ #endif
457
620
  }
458
621
 
459
622
  TEST(TestGpuIndexIVFFlat, AddNaN) {
@@ -462,15 +625,6 @@ TEST(TestGpuIndexIVFFlat, AddNaN) {
462
625
  faiss::gpu::StandardGpuResources res;
463
626
  res.noTempMemory();
464
627
 
465
- faiss::gpu::GpuIndexIVFFlatConfig config;
466
- config.device = opt.device;
467
- config.indicesOptions = opt.indicesOpt;
468
- config.flatConfig.useFloat16 = faiss::gpu::randBool();
469
-
470
- faiss::gpu::GpuIndexIVFFlat gpuIndex(
471
- &res, opt.dim, opt.numCentroids, faiss::METRIC_L2, config);
472
- gpuIndex.nprobe = opt.nprobe;
473
-
474
628
  int numNans = 10;
475
629
  std::vector<float> nans(
476
630
  numNans * opt.dim, std::numeric_limits<float>::quiet_NaN());
@@ -482,6 +636,15 @@ TEST(TestGpuIndexIVFFlat, AddNaN) {
482
636
  }
483
637
 
484
638
  std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
639
+
640
+ faiss::gpu::GpuIndexIVFFlatConfig config;
641
+ config.device = opt.device;
642
+ config.indicesOptions = opt.indicesOpt;
643
+ config.flatConfig.useFloat16 = faiss::gpu::randBool();
644
+ config.use_raft = false;
645
+ faiss::gpu::GpuIndexIVFFlat gpuIndex(
646
+ &res, opt.dim, opt.numCentroids, faiss::METRIC_L2, config);
647
+ gpuIndex.nprobe = opt.nprobe;
485
648
  gpuIndex.train(opt.numTrain, trainVecs.data());
486
649
 
487
650
  // should not crash
@@ -499,6 +662,27 @@ TEST(TestGpuIndexIVFFlat, AddNaN) {
499
662
  opt.k,
500
663
  distance.data(),
501
664
  indices.data());
665
+
666
+ #if defined USE_NVIDIA_RAFT
667
+ config.use_raft = true;
668
+ config.indicesOptions = faiss::gpu::INDICES_64_BIT;
669
+ faiss::gpu::GpuIndexIVFFlat raftGpuIndex(
670
+ &res, opt.dim, opt.numCentroids, faiss::METRIC_L2, config);
671
+ raftGpuIndex.nprobe = opt.nprobe;
672
+ raftGpuIndex.train(opt.numTrain, trainVecs.data());
673
+
674
+ // should not crash
675
+ EXPECT_EQ(raftGpuIndex.ntotal, 0);
676
+ raftGpuIndex.add(numNans, nans.data());
677
+
678
+ // should not crash
679
+ raftGpuIndex.search(
680
+ opt.numQuery,
681
+ queryVecs.data(),
682
+ opt.k,
683
+ distance.data(),
684
+ indices.data());
685
+ #endif
502
686
  }
503
687
 
504
688
  TEST(TestGpuIndexIVFFlat, UnifiedMemory) {
@@ -539,6 +723,7 @@ TEST(TestGpuIndexIVFFlat, UnifiedMemory) {
539
723
  faiss::gpu::GpuIndexIVFFlatConfig config;
540
724
  config.device = device;
541
725
  config.memorySpace = faiss::gpu::MemorySpace::Unified;
726
+ config.use_raft = false;
542
727
 
543
728
  faiss::gpu::GpuIndexIVFFlat gpuIndex(
544
729
  &res, dim, numCentroids, faiss::METRIC_L2, config);
@@ -555,6 +740,26 @@ TEST(TestGpuIndexIVFFlat, UnifiedMemory) {
555
740
  kF32MaxRelErr,
556
741
  0.1f,
557
742
  0.015f);
743
+
744
+ #if defined USE_NVIDIA_RAFT
745
+ config.use_raft = true;
746
+ config.indicesOptions = faiss::gpu::INDICES_64_BIT;
747
+ faiss::gpu::GpuIndexIVFFlat raftGpuIndex(
748
+ &res, dim, numCentroids, faiss::METRIC_L2, config);
749
+ raftGpuIndex.copyFrom(&cpuIndex);
750
+ raftGpuIndex.nprobe = nprobe;
751
+
752
+ faiss::gpu::compareIndices(
753
+ cpuIndex,
754
+ raftGpuIndex,
755
+ numQuery,
756
+ dim,
757
+ k,
758
+ "Unified Memory",
759
+ kF32MaxRelErr,
760
+ 0.1f,
761
+ 0.015f);
762
+ #endif
558
763
  }
559
764
 
560
765
  TEST(TestGpuIndexIVFFlat, LongIVFList) {
@@ -596,6 +801,7 @@ TEST(TestGpuIndexIVFFlat, LongIVFList) {
596
801
 
597
802
  faiss::gpu::GpuIndexIVFFlatConfig config;
598
803
  config.device = device;
804
+ config.use_raft = false;
599
805
 
600
806
  faiss::gpu::GpuIndexIVFFlat gpuIndex(
601
807
  &res, dim, numCentroids, faiss::METRIC_L2, config);
@@ -613,6 +819,92 @@ TEST(TestGpuIndexIVFFlat, LongIVFList) {
613
819
  kF32MaxRelErr,
614
820
  0.1f,
615
821
  0.015f);
822
+
823
+ #if defined USE_NVIDIA_RAFT
824
+ config.use_raft = true;
825
+ config.indicesOptions = faiss::gpu::INDICES_64_BIT;
826
+ faiss::gpu::GpuIndexIVFFlat raftGpuIndex(
827
+ &res, dim, numCentroids, faiss::METRIC_L2, config);
828
+ raftGpuIndex.train(numTrain, trainVecs.data());
829
+ raftGpuIndex.add(numAdd, addVecs.data());
830
+ raftGpuIndex.nprobe = 1;
831
+
832
+ faiss::gpu::compareIndices(
833
+ cpuIndex,
834
+ raftGpuIndex,
835
+ numQuery,
836
+ dim,
837
+ k,
838
+ "Unified Memory",
839
+ kF32MaxRelErr,
840
+ 0.1f,
841
+ 0.015f);
842
+ #endif
843
+ }
844
+
845
+ TEST(TestGpuIndexIVFFlat, Reconstruct_n) {
846
+ Options opt;
847
+
848
+ std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
849
+ std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
850
+
851
+ faiss::IndexFlatL2 cpuQuantizer(opt.dim);
852
+ faiss::IndexIVFFlat cpuIndex(
853
+ &cpuQuantizer, opt.dim, opt.numCentroids, faiss::METRIC_L2);
854
+ cpuIndex.nprobe = opt.nprobe;
855
+ cpuIndex.train(opt.numTrain, trainVecs.data());
856
+ cpuIndex.add(opt.numAdd, addVecs.data());
857
+
858
+ faiss::gpu::StandardGpuResources res;
859
+ res.noTempMemory();
860
+
861
+ faiss::gpu::GpuIndexIVFFlatConfig config;
862
+ config.device = opt.device;
863
+ config.indicesOptions = faiss::gpu::INDICES_64_BIT;
864
+ config.use_raft = false;
865
+
866
+ faiss::gpu::GpuIndexIVFFlat gpuIndex(
867
+ &res, opt.dim, opt.numCentroids, faiss::METRIC_L2, config);
868
+ gpuIndex.nprobe = opt.nprobe;
869
+
870
+ gpuIndex.train(opt.numTrain, trainVecs.data());
871
+ gpuIndex.add(opt.numAdd, addVecs.data());
872
+
873
+ std::vector<float> gpuVals(opt.numAdd * opt.dim);
874
+
875
+ gpuIndex.reconstruct_n(0, gpuIndex.ntotal, gpuVals.data());
876
+
877
+ std::vector<float> cpuVals(opt.numAdd * opt.dim);
878
+
879
+ cpuIndex.reconstruct_n(0, cpuIndex.ntotal, cpuVals.data());
880
+
881
+ EXPECT_EQ(gpuVals, cpuVals);
882
+
883
+ config.indicesOptions = faiss::gpu::INDICES_32_BIT;
884
+
885
+ faiss::gpu::GpuIndexIVFFlat gpuIndex1(
886
+ &res, opt.dim, opt.numCentroids, faiss::METRIC_L2, config);
887
+ gpuIndex1.nprobe = opt.nprobe;
888
+
889
+ gpuIndex1.train(opt.numTrain, trainVecs.data());
890
+ gpuIndex1.add(opt.numAdd, addVecs.data());
891
+
892
+ gpuIndex1.reconstruct_n(0, gpuIndex1.ntotal, gpuVals.data());
893
+
894
+ EXPECT_EQ(gpuVals, cpuVals);
895
+
896
+ config.indicesOptions = faiss::gpu::INDICES_CPU;
897
+
898
+ faiss::gpu::GpuIndexIVFFlat gpuIndex2(
899
+ &res, opt.dim, opt.numCentroids, faiss::METRIC_L2, config);
900
+ gpuIndex2.nprobe = opt.nprobe;
901
+
902
+ gpuIndex2.train(opt.numTrain, trainVecs.data());
903
+ gpuIndex2.add(opt.numAdd, addVecs.data());
904
+
905
+ gpuIndex2.reconstruct_n(0, gpuIndex2.ntotal, gpuVals.data());
906
+
907
+ EXPECT_EQ(gpuVals, cpuVals);
616
908
  }
617
909
 
618
910
  int main(int argc, char** argv) {