faiss 0.1.5 → 0.2.2

Sign up to get free protection for your applications and to get access to all the features.
Files changed (219) hide show
  1. checksums.yaml +4 -4
  2. data/CHANGELOG.md +24 -0
  3. data/README.md +12 -0
  4. data/ext/faiss/ext.cpp +1 -1
  5. data/ext/faiss/extconf.rb +6 -2
  6. data/ext/faiss/index.cpp +114 -43
  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.rb +0 -5
  15. data/lib/faiss/version.rb +1 -1
  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 +24 -10
  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
@@ -5,227 +5,229 @@
5
5
  * LICENSE file in the root directory of this source tree.
6
6
  */
7
7
 
8
-
9
8
  #include <faiss/IndexFlat.h>
10
9
  #include <faiss/IndexScalarQuantizer.h>
11
10
  #include <faiss/gpu/GpuIndexIVFScalarQuantizer.h>
12
11
  #include <faiss/gpu/StandardGpuResources.h>
13
- #include <faiss/gpu/utils/DeviceUtils.h>
14
12
  #include <faiss/gpu/test/TestUtils.h>
15
- #include <cmath>
13
+ #include <faiss/gpu/utils/DeviceUtils.h>
16
14
  #include <gtest/gtest.h>
15
+ #include <cmath>
17
16
  #include <sstream>
18
17
  #include <vector>
19
18
 
20
19
  constexpr float kF32MaxRelErr = 0.03f;
21
20
 
22
21
  struct Options {
23
- Options() {
24
- numAdd = 2 * faiss::gpu::randVal(2000, 5000);
25
- dim = faiss::gpu::randVal(64, 200);
26
-
27
- numCentroids = std::sqrt((float) numAdd / 2);
28
- numTrain = numCentroids * 40;
29
- nprobe = faiss::gpu::randVal(std::min(10, numCentroids), numCentroids);
30
- numQuery = faiss::gpu::randVal(32, 100);
31
-
32
- // Due to the approximate nature of the query and of floating point
33
- // differences between GPU and CPU, to stay within our error bounds, only
34
- // use a small k
35
- k = std::min(faiss::gpu::randVal(10, 30), numAdd / 40);
36
- indicesOpt = faiss::gpu::randSelect({
37
- faiss::gpu::INDICES_CPU,
38
- faiss::gpu::INDICES_32_BIT,
39
- faiss::gpu::INDICES_64_BIT});
40
-
41
- device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
42
- }
43
-
44
- std::string toString() const {
45
- std::stringstream str;
46
- str << "IVFFlat device " << device
47
- << " numVecs " << numAdd
48
- << " dim " << dim
49
- << " numCentroids " << numCentroids
50
- << " nprobe " << nprobe
51
- << " numQuery " << numQuery
52
- << " k " << k
53
- << " indicesOpt " << indicesOpt;
54
-
55
- return str.str();
56
- }
57
-
58
- int numAdd;
59
- int dim;
60
- int numCentroids;
61
- int numTrain;
62
- int nprobe;
63
- int numQuery;
64
- int k;
65
- int device;
66
- faiss::gpu::IndicesOptions indicesOpt;
22
+ Options() {
23
+ numAdd = 2 * faiss::gpu::randVal(2000, 5000);
24
+ dim = faiss::gpu::randVal(64, 200);
25
+
26
+ numCentroids = std::sqrt((float)numAdd / 2);
27
+ numTrain = numCentroids * 40;
28
+ nprobe = faiss::gpu::randVal(std::min(10, numCentroids), numCentroids);
29
+ numQuery = faiss::gpu::randVal(32, 100);
30
+
31
+ // Due to the approximate nature of the query and of floating point
32
+ // differences between GPU and CPU, to stay within our error bounds,
33
+ // only use a small k
34
+ k = std::min(faiss::gpu::randVal(10, 30), numAdd / 40);
35
+ indicesOpt = faiss::gpu::randSelect(
36
+ {faiss::gpu::INDICES_CPU,
37
+ faiss::gpu::INDICES_32_BIT,
38
+ faiss::gpu::INDICES_64_BIT});
39
+
40
+ device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
41
+ }
42
+
43
+ std::string toString() const {
44
+ std::stringstream str;
45
+ str << "IVFFlat device " << device << " numVecs " << numAdd << " dim "
46
+ << dim << " numCentroids " << numCentroids << " nprobe " << nprobe
47
+ << " numQuery " << numQuery << " k " << k << " indicesOpt "
48
+ << indicesOpt;
49
+
50
+ return str.str();
51
+ }
52
+
53
+ int numAdd;
54
+ int dim;
55
+ int numCentroids;
56
+ int numTrain;
57
+ int nprobe;
58
+ int numQuery;
59
+ int k;
60
+ int device;
61
+ faiss::gpu::IndicesOptions indicesOpt;
67
62
  };
68
63
 
69
64
  void runCopyToTest(faiss::ScalarQuantizer::QuantizerType qtype) {
70
- using namespace faiss;
71
- using namespace faiss::gpu;
72
-
73
- Options opt;
74
- std::vector<float> trainVecs = randVecs(opt.numTrain, opt.dim);
75
- std::vector<float> addVecs = randVecs(opt.numAdd, opt.dim);
76
-
77
- StandardGpuResources res;
78
- res.noTempMemory();
79
-
80
- auto config = GpuIndexIVFScalarQuantizerConfig();
81
- config.device = opt.device;
82
-
83
- GpuIndexIVFScalarQuantizer gpuIndex(&res,
84
- opt.dim,
85
- opt.numCentroids,
86
- qtype,
87
- METRIC_L2,
88
- true,
89
- config);
90
- gpuIndex.train(opt.numTrain, trainVecs.data());
91
- gpuIndex.add(opt.numAdd, addVecs.data());
92
- gpuIndex.setNumProbes(opt.nprobe);
93
-
94
- // use garbage values to see if we overwrite then
95
- IndexFlatL2 cpuQuantizer(1);
96
- IndexIVFScalarQuantizer cpuIndex(&cpuQuantizer, 1, 1,
97
- ScalarQuantizer::QuantizerType::QT_6bit,
98
- METRIC_L2);
99
- cpuIndex.nprobe = 1;
100
-
101
- gpuIndex.copyTo(&cpuIndex);
102
-
103
- EXPECT_EQ(cpuIndex.ntotal, gpuIndex.ntotal);
104
- EXPECT_EQ(gpuIndex.ntotal, opt.numAdd);
105
-
106
- EXPECT_EQ(cpuIndex.d, gpuIndex.d);
107
- EXPECT_EQ(cpuIndex.quantizer->d, gpuIndex.quantizer->d);
108
- EXPECT_EQ(cpuIndex.d, opt.dim);
109
- EXPECT_EQ(cpuIndex.nlist, gpuIndex.getNumLists());
110
- EXPECT_EQ(cpuIndex.nprobe, gpuIndex.getNumProbes());
111
-
112
- testIVFEquality(cpuIndex, gpuIndex);
113
-
114
- // Query both objects; results should be equivalent
115
- compareIndices(cpuIndex, gpuIndex,
116
- opt.numQuery, opt.dim, opt.k, opt.toString(),
117
- kF32MaxRelErr,
118
- 0.1f,
119
- 0.015f);
65
+ using namespace faiss;
66
+ using namespace faiss::gpu;
67
+
68
+ Options opt;
69
+ std::vector<float> trainVecs = randVecs(opt.numTrain, opt.dim);
70
+ std::vector<float> addVecs = randVecs(opt.numAdd, opt.dim);
71
+
72
+ StandardGpuResources res;
73
+ res.noTempMemory();
74
+
75
+ auto config = GpuIndexIVFScalarQuantizerConfig();
76
+ config.device = opt.device;
77
+
78
+ GpuIndexIVFScalarQuantizer gpuIndex(
79
+ &res, opt.dim, opt.numCentroids, qtype, METRIC_L2, true, config);
80
+ gpuIndex.train(opt.numTrain, trainVecs.data());
81
+ gpuIndex.add(opt.numAdd, addVecs.data());
82
+ gpuIndex.setNumProbes(opt.nprobe);
83
+
84
+ // use garbage values to see if we overwrite then
85
+ IndexFlatL2 cpuQuantizer(1);
86
+ IndexIVFScalarQuantizer cpuIndex(
87
+ &cpuQuantizer,
88
+ 1,
89
+ 1,
90
+ ScalarQuantizer::QuantizerType::QT_6bit,
91
+ METRIC_L2);
92
+ cpuIndex.nprobe = 1;
93
+
94
+ gpuIndex.copyTo(&cpuIndex);
95
+
96
+ EXPECT_EQ(cpuIndex.ntotal, gpuIndex.ntotal);
97
+ EXPECT_EQ(gpuIndex.ntotal, opt.numAdd);
98
+
99
+ EXPECT_EQ(cpuIndex.d, gpuIndex.d);
100
+ EXPECT_EQ(cpuIndex.quantizer->d, gpuIndex.quantizer->d);
101
+ EXPECT_EQ(cpuIndex.d, opt.dim);
102
+ EXPECT_EQ(cpuIndex.nlist, gpuIndex.getNumLists());
103
+ EXPECT_EQ(cpuIndex.nprobe, gpuIndex.getNumProbes());
104
+
105
+ testIVFEquality(cpuIndex, gpuIndex);
106
+
107
+ // Query both objects; results should be equivalent
108
+ compareIndices(
109
+ cpuIndex,
110
+ gpuIndex,
111
+ opt.numQuery,
112
+ opt.dim,
113
+ opt.k,
114
+ opt.toString(),
115
+ kF32MaxRelErr,
116
+ 0.1f,
117
+ 0.015f);
120
118
  }
121
119
 
122
120
  TEST(TestGpuIndexIVFScalarQuantizer, CopyTo_fp16) {
123
- runCopyToTest(faiss::ScalarQuantizer::QuantizerType::QT_fp16);
121
+ runCopyToTest(faiss::ScalarQuantizer::QuantizerType::QT_fp16);
124
122
  }
125
123
 
126
124
  TEST(TestGpuIndexIVFScalarQuantizer, CopyTo_8bit) {
127
- runCopyToTest(faiss::ScalarQuantizer::QuantizerType::QT_8bit);
125
+ runCopyToTest(faiss::ScalarQuantizer::QuantizerType::QT_8bit);
128
126
  }
129
127
 
130
128
  TEST(TestGpuIndexIVFScalarQuantizer, CopyTo_8bit_uniform) {
131
- runCopyToTest(faiss::ScalarQuantizer::QuantizerType::QT_8bit_uniform);
129
+ runCopyToTest(faiss::ScalarQuantizer::QuantizerType::QT_8bit_uniform);
132
130
  }
133
131
 
134
132
  TEST(TestGpuIndexIVFScalarQuantizer, CopyTo_6bit) {
135
- runCopyToTest(faiss::ScalarQuantizer::QuantizerType::QT_6bit);
133
+ runCopyToTest(faiss::ScalarQuantizer::QuantizerType::QT_6bit);
136
134
  }
137
135
 
138
136
  TEST(TestGpuIndexIVFScalarQuantizer, CopyTo_4bit) {
139
- runCopyToTest(faiss::ScalarQuantizer::QuantizerType::QT_4bit);
137
+ runCopyToTest(faiss::ScalarQuantizer::QuantizerType::QT_4bit);
140
138
  }
141
139
 
142
140
  TEST(TestGpuIndexIVFScalarQuantizer, CopyTo_4bit_uniform) {
143
- runCopyToTest(faiss::ScalarQuantizer::QuantizerType::QT_4bit_uniform);
141
+ runCopyToTest(faiss::ScalarQuantizer::QuantizerType::QT_4bit_uniform);
144
142
  }
145
143
 
146
144
  void runCopyFromTest(faiss::ScalarQuantizer::QuantizerType qtype) {
147
- using namespace faiss;
148
- using namespace faiss::gpu;
149
-
150
- Options opt;
151
- std::vector<float> trainVecs = randVecs(opt.numTrain, opt.dim);
152
- std::vector<float> addVecs = randVecs(opt.numAdd, opt.dim);
153
-
154
- IndexFlatL2 cpuQuantizer(opt.dim);
155
- IndexIVFScalarQuantizer cpuIndex(&cpuQuantizer, opt.dim, opt.numCentroids,
156
- qtype,
157
- METRIC_L2);
158
-
159
- cpuIndex.nprobe = opt.nprobe;
160
- cpuIndex.train(opt.numTrain, trainVecs.data());
161
- cpuIndex.add(opt.numAdd, addVecs.data());
162
-
163
- // use garbage values to see if we overwrite then
164
- StandardGpuResources res;
165
- res.noTempMemory();
166
-
167
- auto config = GpuIndexIVFScalarQuantizerConfig();
168
- config.device = opt.device;
169
-
170
- GpuIndexIVFScalarQuantizer gpuIndex(
171
- &res,
172
- 1,
173
- 1,
174
- ScalarQuantizer::QuantizerType::QT_4bit,
175
- METRIC_L2,
176
- false,
177
- config);
178
- gpuIndex.setNumProbes(1);
179
-
180
- gpuIndex.copyFrom(&cpuIndex);
181
-
182
- EXPECT_EQ(cpuIndex.ntotal, gpuIndex.ntotal);
183
- EXPECT_EQ(gpuIndex.ntotal, opt.numAdd);
184
-
185
- EXPECT_EQ(cpuIndex.d, gpuIndex.d);
186
- EXPECT_EQ(cpuIndex.d, opt.dim);
187
- EXPECT_EQ(cpuIndex.nlist, gpuIndex.getNumLists());
188
- EXPECT_EQ(cpuIndex.nprobe, gpuIndex.getNumProbes());
189
-
190
- testIVFEquality(cpuIndex, gpuIndex);
191
-
192
- // Query both objects; results should be equivalent
193
- compareIndices(cpuIndex, gpuIndex,
194
- opt.numQuery, opt.dim, opt.k, opt.toString(),
195
- kF32MaxRelErr,
196
- 0.1f,
197
- 0.015f);
145
+ using namespace faiss;
146
+ using namespace faiss::gpu;
147
+
148
+ Options opt;
149
+ std::vector<float> trainVecs = randVecs(opt.numTrain, opt.dim);
150
+ std::vector<float> addVecs = randVecs(opt.numAdd, opt.dim);
151
+
152
+ IndexFlatL2 cpuQuantizer(opt.dim);
153
+ IndexIVFScalarQuantizer cpuIndex(
154
+ &cpuQuantizer, opt.dim, opt.numCentroids, qtype, METRIC_L2);
155
+
156
+ cpuIndex.nprobe = opt.nprobe;
157
+ cpuIndex.train(opt.numTrain, trainVecs.data());
158
+ cpuIndex.add(opt.numAdd, addVecs.data());
159
+
160
+ // use garbage values to see if we overwrite then
161
+ StandardGpuResources res;
162
+ res.noTempMemory();
163
+
164
+ auto config = GpuIndexIVFScalarQuantizerConfig();
165
+ config.device = opt.device;
166
+
167
+ GpuIndexIVFScalarQuantizer gpuIndex(
168
+ &res,
169
+ 1,
170
+ 1,
171
+ ScalarQuantizer::QuantizerType::QT_4bit,
172
+ METRIC_L2,
173
+ false,
174
+ config);
175
+ gpuIndex.setNumProbes(1);
176
+
177
+ gpuIndex.copyFrom(&cpuIndex);
178
+
179
+ EXPECT_EQ(cpuIndex.ntotal, gpuIndex.ntotal);
180
+ EXPECT_EQ(gpuIndex.ntotal, opt.numAdd);
181
+
182
+ EXPECT_EQ(cpuIndex.d, gpuIndex.d);
183
+ EXPECT_EQ(cpuIndex.d, opt.dim);
184
+ EXPECT_EQ(cpuIndex.nlist, gpuIndex.getNumLists());
185
+ EXPECT_EQ(cpuIndex.nprobe, gpuIndex.getNumProbes());
186
+
187
+ testIVFEquality(cpuIndex, gpuIndex);
188
+
189
+ // Query both objects; results should be equivalent
190
+ compareIndices(
191
+ cpuIndex,
192
+ gpuIndex,
193
+ opt.numQuery,
194
+ opt.dim,
195
+ opt.k,
196
+ opt.toString(),
197
+ kF32MaxRelErr,
198
+ 0.1f,
199
+ 0.015f);
198
200
  }
199
201
 
200
202
  TEST(TestGpuIndexIVFScalarQuantizer, CopyFrom_fp16) {
201
- runCopyFromTest(faiss::ScalarQuantizer::QuantizerType::QT_fp16);
203
+ runCopyFromTest(faiss::ScalarQuantizer::QuantizerType::QT_fp16);
202
204
  }
203
205
 
204
206
  TEST(TestGpuIndexIVFScalarQuantizer, CopyFrom_8bit) {
205
- runCopyFromTest(faiss::ScalarQuantizer::QuantizerType::QT_8bit);
207
+ runCopyFromTest(faiss::ScalarQuantizer::QuantizerType::QT_8bit);
206
208
  }
207
209
 
208
210
  TEST(TestGpuIndexIVFScalarQuantizer, CopyFrom_8bit_uniform) {
209
- runCopyFromTest(faiss::ScalarQuantizer::QuantizerType::QT_8bit_uniform);
211
+ runCopyFromTest(faiss::ScalarQuantizer::QuantizerType::QT_8bit_uniform);
210
212
  }
211
213
 
212
214
  TEST(TestGpuIndexIVFScalarQuantizer, CopyFrom_6bit) {
213
- runCopyFromTest(faiss::ScalarQuantizer::QuantizerType::QT_6bit);
215
+ runCopyFromTest(faiss::ScalarQuantizer::QuantizerType::QT_6bit);
214
216
  }
215
217
 
216
218
  TEST(TestGpuIndexIVFScalarQuantizer, CopyFrom_4bit) {
217
- runCopyFromTest(faiss::ScalarQuantizer::QuantizerType::QT_4bit);
219
+ runCopyFromTest(faiss::ScalarQuantizer::QuantizerType::QT_4bit);
218
220
  }
219
221
 
220
222
  TEST(TestGpuIndexIVFScalarQuantizer, CopyFrom_4bit_uniform) {
221
- runCopyFromTest(faiss::ScalarQuantizer::QuantizerType::QT_4bit_uniform);
223
+ runCopyFromTest(faiss::ScalarQuantizer::QuantizerType::QT_4bit_uniform);
222
224
  }
223
225
 
224
226
  int main(int argc, char** argv) {
225
- testing::InitGoogleTest(&argc, argv);
227
+ testing::InitGoogleTest(&argc, argv);
226
228
 
227
- // just run with a fixed test seed
228
- faiss::gpu::setTestSeed(100);
229
+ // just run with a fixed test seed
230
+ faiss::gpu::setTestSeed(100);
229
231
 
230
- return RUN_ALL_TESTS();
232
+ return RUN_ALL_TESTS();
231
233
  }
@@ -5,80 +5,88 @@
5
5
  * LICENSE file in the root directory of this source tree.
6
6
  */
7
7
 
8
-
9
8
  #include <faiss/IndexFlat.h>
10
9
  #include <faiss/gpu/GpuIndexFlat.h>
11
10
  #include <faiss/gpu/StandardGpuResources.h>
12
- #include <faiss/gpu/utils/DeviceUtils.h>
13
11
  #include <faiss/gpu/test/TestUtils.h>
12
+ #include <faiss/gpu/utils/DeviceUtils.h>
14
13
  #include <gtest/gtest.h>
15
14
 
16
15
  // Test to see if we can recover after attempting to allocate too much GPU
17
16
  // memory
18
17
  TEST(TestGpuMemoryException, AddException) {
19
- size_t numBrokenAdd = std::numeric_limits<int>::max();
20
- size_t numRealAdd = 10000;
21
- size_t devFree = 0;
22
- size_t devTotal = 0;
18
+ size_t numBrokenAdd = std::numeric_limits<int>::max();
19
+ size_t numRealAdd = 10000;
20
+ size_t devFree = 0;
21
+ size_t devTotal = 0;
23
22
 
24
- CUDA_VERIFY(cudaMemGetInfo(&devFree, &devTotal));
23
+ CUDA_VERIFY(cudaMemGetInfo(&devFree, &devTotal));
25
24
 
26
- // Figure out the dimensionality needed to get at least greater than devTotal
27
- size_t brokenAddDims = ((devTotal / sizeof(float)) / numBrokenAdd) + 1;
28
- size_t realAddDims = 128;
25
+ // Figure out the dimensionality needed to get at least greater than
26
+ // devTotal
27
+ size_t brokenAddDims = ((devTotal / sizeof(float)) / numBrokenAdd) + 1;
28
+ size_t realAddDims = 128;
29
29
 
30
- faiss::gpu::StandardGpuResources res;
30
+ faiss::gpu::StandardGpuResources res;
31
31
 
32
- faiss::gpu::GpuIndexFlatConfig config;
33
- config.device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
32
+ faiss::gpu::GpuIndexFlatConfig config;
33
+ config.device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
34
34
 
35
- faiss::gpu::GpuIndexFlatL2
36
- gpuIndexL2Broken(&res, (int) brokenAddDims, config);
37
- faiss::gpu::GpuIndexFlatL2
38
- gpuIndexL2(&res, (int) realAddDims, config);
39
- faiss::IndexFlatL2
40
- cpuIndex((int) realAddDims);
35
+ faiss::gpu::GpuIndexFlatL2 gpuIndexL2Broken(
36
+ &res, (int)brokenAddDims, config);
37
+ faiss::gpu::GpuIndexFlatL2 gpuIndexL2(&res, (int)realAddDims, config);
38
+ faiss::IndexFlatL2 cpuIndex((int)realAddDims);
41
39
 
42
- // Should throw on attempting to allocate too much data
43
- {
44
- // allocate memory without initialization
45
- auto vecs =
46
- std::unique_ptr<float[]>(new float[numBrokenAdd * brokenAddDims]);
47
- EXPECT_THROW(gpuIndexL2Broken.add(numBrokenAdd, vecs.get()),
48
- faiss::FaissException);
49
- }
40
+ // Should throw on attempting to allocate too much data
41
+ {
42
+ // allocate memory without initialization
43
+ auto vecs = std::unique_ptr<float[]>(
44
+ new float[numBrokenAdd * brokenAddDims]);
45
+ EXPECT_THROW(
46
+ gpuIndexL2Broken.add(numBrokenAdd, vecs.get()),
47
+ faiss::FaissException);
48
+ }
50
49
 
51
- // Should be able to add a smaller set of data now
52
- {
53
- auto vecs = faiss::gpu::randVecs(numRealAdd, realAddDims);
54
- EXPECT_NO_THROW(gpuIndexL2.add(numRealAdd, vecs.data()));
55
- cpuIndex.add(numRealAdd, vecs.data());
56
- }
50
+ // Should be able to add a smaller set of data now
51
+ {
52
+ auto vecs = faiss::gpu::randVecs(numRealAdd, realAddDims);
53
+ EXPECT_NO_THROW(gpuIndexL2.add(numRealAdd, vecs.data()));
54
+ cpuIndex.add(numRealAdd, vecs.data());
55
+ }
57
56
 
58
- // Should throw on attempting to allocate too much data
59
- {
60
- // allocate memory without initialization
61
- auto vecs =
62
- std::unique_ptr<float[]>(new float[numBrokenAdd * brokenAddDims]);
63
- EXPECT_THROW(gpuIndexL2Broken.add(numBrokenAdd, vecs.get()),
64
- faiss::FaissException);
65
- }
57
+ // Should throw on attempting to allocate too much data
58
+ {
59
+ // allocate memory without initialization
60
+ auto vecs = std::unique_ptr<float[]>(
61
+ new float[numBrokenAdd * brokenAddDims]);
62
+ EXPECT_THROW(
63
+ gpuIndexL2Broken.add(numBrokenAdd, vecs.get()),
64
+ faiss::FaissException);
65
+ }
66
66
 
67
- // Should be able to query results from what we had before
68
- {
69
- size_t numQuery = 10;
70
- auto vecs = faiss::gpu::randVecs(numQuery, realAddDims);
71
- EXPECT_NO_THROW(compareIndices(vecs, cpuIndex, gpuIndexL2,
72
- numQuery, realAddDims, 50, "",
73
- 6e-3f, 0.1f, 0.015f));
74
- }
67
+ // Should be able to query results from what we had before
68
+ {
69
+ size_t numQuery = 10;
70
+ auto vecs = faiss::gpu::randVecs(numQuery, realAddDims);
71
+ EXPECT_NO_THROW(compareIndices(
72
+ vecs,
73
+ cpuIndex,
74
+ gpuIndexL2,
75
+ numQuery,
76
+ realAddDims,
77
+ 50,
78
+ "",
79
+ 6e-3f,
80
+ 0.1f,
81
+ 0.015f));
82
+ }
75
83
  }
76
84
 
77
85
  int main(int argc, char** argv) {
78
- testing::InitGoogleTest(&argc, argv);
86
+ testing::InitGoogleTest(&argc, argv);
79
87
 
80
- // just run with a fixed test seed
81
- faiss::gpu::setTestSeed(100);
88
+ // just run with a fixed test seed
89
+ faiss::gpu::setTestSeed(100);
82
90
 
83
- return RUN_ALL_TESTS();
91
+ return RUN_ALL_TESTS();
84
92
  }