faiss 0.1.7 → 0.2.3

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 +18 -0
  3. data/README.md +7 -7
  4. data/ext/faiss/ext.cpp +1 -1
  5. data/ext/faiss/extconf.rb +8 -2
  6. data/ext/faiss/index.cpp +102 -69
  7. data/ext/faiss/index_binary.cpp +24 -30
  8. data/ext/faiss/kmeans.cpp +20 -16
  9. data/ext/faiss/numo.hpp +867 -0
  10. data/ext/faiss/pca_matrix.cpp +13 -14
  11. data/ext/faiss/product_quantizer.cpp +23 -24
  12. data/ext/faiss/utils.cpp +10 -37
  13. data/ext/faiss/utils.h +2 -13
  14. data/lib/faiss/version.rb +1 -1
  15. data/lib/faiss.rb +0 -5
  16. data/vendor/faiss/faiss/AutoTune.cpp +292 -291
  17. data/vendor/faiss/faiss/AutoTune.h +55 -56
  18. data/vendor/faiss/faiss/Clustering.cpp +334 -195
  19. data/vendor/faiss/faiss/Clustering.h +88 -35
  20. data/vendor/faiss/faiss/IVFlib.cpp +171 -195
  21. data/vendor/faiss/faiss/IVFlib.h +48 -51
  22. data/vendor/faiss/faiss/Index.cpp +85 -103
  23. data/vendor/faiss/faiss/Index.h +54 -48
  24. data/vendor/faiss/faiss/Index2Layer.cpp +139 -164
  25. data/vendor/faiss/faiss/Index2Layer.h +22 -22
  26. data/vendor/faiss/faiss/IndexBinary.cpp +45 -37
  27. data/vendor/faiss/faiss/IndexBinary.h +140 -132
  28. data/vendor/faiss/faiss/IndexBinaryFlat.cpp +73 -53
  29. data/vendor/faiss/faiss/IndexBinaryFlat.h +29 -24
  30. data/vendor/faiss/faiss/IndexBinaryFromFloat.cpp +46 -43
  31. data/vendor/faiss/faiss/IndexBinaryFromFloat.h +16 -15
  32. data/vendor/faiss/faiss/IndexBinaryHNSW.cpp +215 -232
  33. data/vendor/faiss/faiss/IndexBinaryHNSW.h +25 -24
  34. data/vendor/faiss/faiss/IndexBinaryHash.cpp +182 -177
  35. data/vendor/faiss/faiss/IndexBinaryHash.h +41 -34
  36. data/vendor/faiss/faiss/IndexBinaryIVF.cpp +489 -461
  37. data/vendor/faiss/faiss/IndexBinaryIVF.h +97 -68
  38. data/vendor/faiss/faiss/IndexFlat.cpp +116 -147
  39. data/vendor/faiss/faiss/IndexFlat.h +35 -46
  40. data/vendor/faiss/faiss/IndexHNSW.cpp +372 -348
  41. data/vendor/faiss/faiss/IndexHNSW.h +57 -41
  42. data/vendor/faiss/faiss/IndexIVF.cpp +474 -454
  43. data/vendor/faiss/faiss/IndexIVF.h +146 -113
  44. data/vendor/faiss/faiss/IndexIVFFlat.cpp +248 -250
  45. data/vendor/faiss/faiss/IndexIVFFlat.h +48 -51
  46. data/vendor/faiss/faiss/IndexIVFPQ.cpp +457 -516
  47. data/vendor/faiss/faiss/IndexIVFPQ.h +74 -66
  48. data/vendor/faiss/faiss/IndexIVFPQFastScan.cpp +406 -372
  49. data/vendor/faiss/faiss/IndexIVFPQFastScan.h +82 -57
  50. data/vendor/faiss/faiss/IndexIVFPQR.cpp +104 -102
  51. data/vendor/faiss/faiss/IndexIVFPQR.h +33 -28
  52. data/vendor/faiss/faiss/IndexIVFSpectralHash.cpp +125 -133
  53. data/vendor/faiss/faiss/IndexIVFSpectralHash.h +19 -21
  54. data/vendor/faiss/faiss/IndexLSH.cpp +75 -96
  55. data/vendor/faiss/faiss/IndexLSH.h +21 -26
  56. data/vendor/faiss/faiss/IndexLattice.cpp +42 -56
  57. data/vendor/faiss/faiss/IndexLattice.h +11 -16
  58. data/vendor/faiss/faiss/IndexNNDescent.cpp +231 -0
  59. data/vendor/faiss/faiss/IndexNNDescent.h +72 -0
  60. data/vendor/faiss/faiss/IndexNSG.cpp +303 -0
  61. data/vendor/faiss/faiss/IndexNSG.h +85 -0
  62. data/vendor/faiss/faiss/IndexPQ.cpp +405 -464
  63. data/vendor/faiss/faiss/IndexPQ.h +64 -67
  64. data/vendor/faiss/faiss/IndexPQFastScan.cpp +143 -170
  65. data/vendor/faiss/faiss/IndexPQFastScan.h +46 -32
  66. data/vendor/faiss/faiss/IndexPreTransform.cpp +120 -150
  67. data/vendor/faiss/faiss/IndexPreTransform.h +33 -36
  68. data/vendor/faiss/faiss/IndexRefine.cpp +115 -131
  69. data/vendor/faiss/faiss/IndexRefine.h +22 -23
  70. data/vendor/faiss/faiss/IndexReplicas.cpp +147 -153
  71. data/vendor/faiss/faiss/IndexReplicas.h +62 -56
  72. data/vendor/faiss/faiss/IndexResidual.cpp +291 -0
  73. data/vendor/faiss/faiss/IndexResidual.h +152 -0
  74. data/vendor/faiss/faiss/IndexScalarQuantizer.cpp +120 -155
  75. data/vendor/faiss/faiss/IndexScalarQuantizer.h +41 -45
  76. data/vendor/faiss/faiss/IndexShards.cpp +256 -240
  77. data/vendor/faiss/faiss/IndexShards.h +85 -73
  78. data/vendor/faiss/faiss/MatrixStats.cpp +112 -97
  79. data/vendor/faiss/faiss/MatrixStats.h +7 -10
  80. data/vendor/faiss/faiss/MetaIndexes.cpp +135 -157
  81. data/vendor/faiss/faiss/MetaIndexes.h +40 -34
  82. data/vendor/faiss/faiss/MetricType.h +7 -7
  83. data/vendor/faiss/faiss/VectorTransform.cpp +652 -474
  84. data/vendor/faiss/faiss/VectorTransform.h +61 -89
  85. data/vendor/faiss/faiss/clone_index.cpp +77 -73
  86. data/vendor/faiss/faiss/clone_index.h +4 -9
  87. data/vendor/faiss/faiss/gpu/GpuAutoTune.cpp +33 -38
  88. data/vendor/faiss/faiss/gpu/GpuAutoTune.h +11 -9
  89. data/vendor/faiss/faiss/gpu/GpuCloner.cpp +197 -170
  90. data/vendor/faiss/faiss/gpu/GpuCloner.h +53 -35
  91. data/vendor/faiss/faiss/gpu/GpuClonerOptions.cpp +12 -14
  92. data/vendor/faiss/faiss/gpu/GpuClonerOptions.h +27 -25
  93. data/vendor/faiss/faiss/gpu/GpuDistance.h +116 -112
  94. data/vendor/faiss/faiss/gpu/GpuFaissAssert.h +1 -2
  95. data/vendor/faiss/faiss/gpu/GpuIndex.h +134 -137
  96. data/vendor/faiss/faiss/gpu/GpuIndexBinaryFlat.h +76 -73
  97. data/vendor/faiss/faiss/gpu/GpuIndexFlat.h +173 -162
  98. data/vendor/faiss/faiss/gpu/GpuIndexIVF.h +67 -64
  99. data/vendor/faiss/faiss/gpu/GpuIndexIVFFlat.h +89 -86
  100. data/vendor/faiss/faiss/gpu/GpuIndexIVFPQ.h +150 -141
  101. data/vendor/faiss/faiss/gpu/GpuIndexIVFScalarQuantizer.h +101 -103
  102. data/vendor/faiss/faiss/gpu/GpuIndicesOptions.h +17 -16
  103. data/vendor/faiss/faiss/gpu/GpuResources.cpp +116 -128
  104. data/vendor/faiss/faiss/gpu/GpuResources.h +182 -186
  105. data/vendor/faiss/faiss/gpu/StandardGpuResources.cpp +433 -422
  106. data/vendor/faiss/faiss/gpu/StandardGpuResources.h +131 -130
  107. data/vendor/faiss/faiss/gpu/impl/InterleavedCodes.cpp +468 -456
  108. data/vendor/faiss/faiss/gpu/impl/InterleavedCodes.h +25 -19
  109. data/vendor/faiss/faiss/gpu/impl/RemapIndices.cpp +22 -20
  110. data/vendor/faiss/faiss/gpu/impl/RemapIndices.h +9 -8
  111. data/vendor/faiss/faiss/gpu/perf/IndexWrapper-inl.h +39 -44
  112. data/vendor/faiss/faiss/gpu/perf/IndexWrapper.h +16 -14
  113. data/vendor/faiss/faiss/gpu/perf/PerfClustering.cpp +77 -71
  114. data/vendor/faiss/faiss/gpu/perf/PerfIVFPQAdd.cpp +109 -88
  115. data/vendor/faiss/faiss/gpu/perf/WriteIndex.cpp +75 -64
  116. data/vendor/faiss/faiss/gpu/test/TestCodePacking.cpp +230 -215
  117. data/vendor/faiss/faiss/gpu/test/TestGpuIndexBinaryFlat.cpp +80 -86
  118. data/vendor/faiss/faiss/gpu/test/TestGpuIndexFlat.cpp +284 -277
  119. data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFFlat.cpp +416 -416
  120. data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFPQ.cpp +611 -517
  121. data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFScalarQuantizer.cpp +166 -164
  122. data/vendor/faiss/faiss/gpu/test/TestGpuMemoryException.cpp +61 -53
  123. data/vendor/faiss/faiss/gpu/test/TestUtils.cpp +274 -238
  124. data/vendor/faiss/faiss/gpu/test/TestUtils.h +73 -57
  125. data/vendor/faiss/faiss/gpu/test/demo_ivfpq_indexing_gpu.cpp +47 -50
  126. data/vendor/faiss/faiss/gpu/utils/DeviceUtils.h +79 -72
  127. data/vendor/faiss/faiss/gpu/utils/StackDeviceMemory.cpp +140 -146
  128. data/vendor/faiss/faiss/gpu/utils/StackDeviceMemory.h +69 -71
  129. data/vendor/faiss/faiss/gpu/utils/StaticUtils.h +21 -16
  130. data/vendor/faiss/faiss/gpu/utils/Timer.cpp +25 -29
  131. data/vendor/faiss/faiss/gpu/utils/Timer.h +30 -29
  132. data/vendor/faiss/faiss/impl/AdditiveQuantizer.cpp +270 -0
  133. data/vendor/faiss/faiss/impl/AdditiveQuantizer.h +115 -0
  134. data/vendor/faiss/faiss/impl/AuxIndexStructures.cpp +90 -120
  135. data/vendor/faiss/faiss/impl/AuxIndexStructures.h +81 -65
  136. data/vendor/faiss/faiss/impl/FaissAssert.h +73 -58
  137. data/vendor/faiss/faiss/impl/FaissException.cpp +56 -48
  138. data/vendor/faiss/faiss/impl/FaissException.h +41 -29
  139. data/vendor/faiss/faiss/impl/HNSW.cpp +595 -611
  140. data/vendor/faiss/faiss/impl/HNSW.h +179 -200
  141. data/vendor/faiss/faiss/impl/LocalSearchQuantizer.cpp +672 -0
  142. data/vendor/faiss/faiss/impl/LocalSearchQuantizer.h +172 -0
  143. data/vendor/faiss/faiss/impl/NNDescent.cpp +487 -0
  144. data/vendor/faiss/faiss/impl/NNDescent.h +154 -0
  145. data/vendor/faiss/faiss/impl/NSG.cpp +682 -0
  146. data/vendor/faiss/faiss/impl/NSG.h +199 -0
  147. data/vendor/faiss/faiss/impl/PolysemousTraining.cpp +484 -454
  148. data/vendor/faiss/faiss/impl/PolysemousTraining.h +52 -55
  149. data/vendor/faiss/faiss/impl/ProductQuantizer-inl.h +26 -47
  150. data/vendor/faiss/faiss/impl/ProductQuantizer.cpp +469 -459
  151. data/vendor/faiss/faiss/impl/ProductQuantizer.h +76 -87
  152. data/vendor/faiss/faiss/impl/ResidualQuantizer.cpp +448 -0
  153. data/vendor/faiss/faiss/impl/ResidualQuantizer.h +130 -0
  154. data/vendor/faiss/faiss/impl/ResultHandler.h +96 -132
  155. data/vendor/faiss/faiss/impl/ScalarQuantizer.cpp +648 -701
  156. data/vendor/faiss/faiss/impl/ScalarQuantizer.h +48 -46
  157. data/vendor/faiss/faiss/impl/ThreadedIndex-inl.h +129 -131
  158. data/vendor/faiss/faiss/impl/ThreadedIndex.h +61 -55
  159. data/vendor/faiss/faiss/impl/index_read.cpp +547 -479
  160. data/vendor/faiss/faiss/impl/index_write.cpp +497 -407
  161. data/vendor/faiss/faiss/impl/io.cpp +75 -94
  162. data/vendor/faiss/faiss/impl/io.h +31 -41
  163. data/vendor/faiss/faiss/impl/io_macros.h +40 -29
  164. data/vendor/faiss/faiss/impl/lattice_Zn.cpp +137 -186
  165. data/vendor/faiss/faiss/impl/lattice_Zn.h +40 -51
  166. data/vendor/faiss/faiss/impl/platform_macros.h +29 -8
  167. data/vendor/faiss/faiss/impl/pq4_fast_scan.cpp +77 -124
  168. data/vendor/faiss/faiss/impl/pq4_fast_scan.h +39 -48
  169. data/vendor/faiss/faiss/impl/pq4_fast_scan_search_1.cpp +41 -52
  170. data/vendor/faiss/faiss/impl/pq4_fast_scan_search_qbs.cpp +80 -117
  171. data/vendor/faiss/faiss/impl/simd_result_handlers.h +109 -137
  172. data/vendor/faiss/faiss/index_factory.cpp +269 -218
  173. data/vendor/faiss/faiss/index_factory.h +6 -7
  174. data/vendor/faiss/faiss/index_io.h +23 -26
  175. data/vendor/faiss/faiss/invlists/BlockInvertedLists.cpp +67 -75
  176. data/vendor/faiss/faiss/invlists/BlockInvertedLists.h +22 -24
  177. data/vendor/faiss/faiss/invlists/DirectMap.cpp +96 -112
  178. data/vendor/faiss/faiss/invlists/DirectMap.h +29 -33
  179. data/vendor/faiss/faiss/invlists/InvertedLists.cpp +307 -364
  180. data/vendor/faiss/faiss/invlists/InvertedLists.h +151 -151
  181. data/vendor/faiss/faiss/invlists/InvertedListsIOHook.cpp +29 -34
  182. data/vendor/faiss/faiss/invlists/InvertedListsIOHook.h +17 -18
  183. data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.cpp +257 -293
  184. data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.h +50 -45
  185. data/vendor/faiss/faiss/python/python_callbacks.cpp +23 -26
  186. data/vendor/faiss/faiss/python/python_callbacks.h +9 -16
  187. data/vendor/faiss/faiss/utils/AlignedTable.h +79 -44
  188. data/vendor/faiss/faiss/utils/Heap.cpp +40 -48
  189. data/vendor/faiss/faiss/utils/Heap.h +186 -209
  190. data/vendor/faiss/faiss/utils/WorkerThread.cpp +67 -76
  191. data/vendor/faiss/faiss/utils/WorkerThread.h +32 -33
  192. data/vendor/faiss/faiss/utils/distances.cpp +301 -310
  193. data/vendor/faiss/faiss/utils/distances.h +133 -118
  194. data/vendor/faiss/faiss/utils/distances_simd.cpp +456 -516
  195. data/vendor/faiss/faiss/utils/extra_distances-inl.h +117 -0
  196. data/vendor/faiss/faiss/utils/extra_distances.cpp +113 -232
  197. data/vendor/faiss/faiss/utils/extra_distances.h +30 -29
  198. data/vendor/faiss/faiss/utils/hamming-inl.h +260 -209
  199. data/vendor/faiss/faiss/utils/hamming.cpp +375 -469
  200. data/vendor/faiss/faiss/utils/hamming.h +62 -85
  201. data/vendor/faiss/faiss/utils/ordered_key_value.h +16 -18
  202. data/vendor/faiss/faiss/utils/partitioning.cpp +393 -318
  203. data/vendor/faiss/faiss/utils/partitioning.h +26 -21
  204. data/vendor/faiss/faiss/utils/quantize_lut.cpp +78 -66
  205. data/vendor/faiss/faiss/utils/quantize_lut.h +22 -20
  206. data/vendor/faiss/faiss/utils/random.cpp +39 -63
  207. data/vendor/faiss/faiss/utils/random.h +13 -16
  208. data/vendor/faiss/faiss/utils/simdlib.h +4 -2
  209. data/vendor/faiss/faiss/utils/simdlib_avx2.h +88 -85
  210. data/vendor/faiss/faiss/utils/simdlib_emulated.h +226 -165
  211. data/vendor/faiss/faiss/utils/simdlib_neon.h +832 -0
  212. data/vendor/faiss/faiss/utils/utils.cpp +304 -287
  213. data/vendor/faiss/faiss/utils/utils.h +53 -48
  214. metadata +26 -12
  215. data/lib/faiss/index.rb +0 -20
  216. data/lib/faiss/index_binary.rb +0 -20
  217. data/lib/faiss/kmeans.rb +0 -15
  218. data/lib/faiss/pca_matrix.rb +0 -15
  219. data/lib/faiss/product_quantizer.rb +0 -22
@@ -5,15 +5,14 @@
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/IndexIVFFlat.h>
11
10
  #include <faiss/gpu/GpuIndexIVFFlat.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
 
@@ -21,76 +20,157 @@
21
20
  constexpr float kF16MaxRelErr = 0.3f;
22
21
  constexpr float kF32MaxRelErr = 0.03f;
23
22
 
24
-
25
23
  struct Options {
26
- Options() {
27
- numAdd = 2 * faiss::gpu::randVal(2000, 5000);
28
- dim = faiss::gpu::randVal(64, 200);
29
-
30
- numCentroids = std::sqrt((float) numAdd / 2);
31
- numTrain = numCentroids * 40;
32
- nprobe = faiss::gpu::randVal(std::min(10, numCentroids), numCentroids);
33
- numQuery = faiss::gpu::randVal(32, 100);
34
-
35
- // Due to the approximate nature of the query and of floating point
36
- // differences between GPU and CPU, to stay within our error bounds, only
37
- // use a small k
38
- k = std::min(faiss::gpu::randVal(10, 30), numAdd / 40);
39
- indicesOpt = faiss::gpu::randSelect({
40
- faiss::gpu::INDICES_CPU,
41
- faiss::gpu::INDICES_32_BIT,
42
- faiss::gpu::INDICES_64_BIT});
43
-
44
- device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
45
- }
46
-
47
- std::string toString() const {
48
- std::stringstream str;
49
- str << "IVFFlat device " << device
50
- << " numVecs " << numAdd
51
- << " dim " << dim
52
- << " numCentroids " << numCentroids
53
- << " nprobe " << nprobe
54
- << " numQuery " << numQuery
55
- << " k " << k
56
- << " indicesOpt " << indicesOpt;
57
-
58
- return str.str();
59
- }
60
-
61
- int numAdd;
62
- int dim;
63
- int numCentroids;
64
- int numTrain;
65
- int nprobe;
66
- int numQuery;
67
- int k;
68
- int device;
69
- faiss::gpu::IndicesOptions indicesOpt;
24
+ Options() {
25
+ numAdd = 2 * faiss::gpu::randVal(2000, 5000);
26
+ dim = faiss::gpu::randVal(64, 200);
27
+
28
+ numCentroids = std::sqrt((float)numAdd / 2);
29
+ numTrain = numCentroids * 40;
30
+ nprobe = faiss::gpu::randVal(std::min(10, numCentroids), numCentroids);
31
+ numQuery = faiss::gpu::randVal(32, 100);
32
+
33
+ // Due to the approximate nature of the query and of floating point
34
+ // differences between GPU and CPU, to stay within our error bounds,
35
+ // only use a small k
36
+ k = std::min(faiss::gpu::randVal(10, 30), numAdd / 40);
37
+ indicesOpt = faiss::gpu::randSelect(
38
+ {faiss::gpu::INDICES_CPU,
39
+ faiss::gpu::INDICES_32_BIT,
40
+ faiss::gpu::INDICES_64_BIT});
41
+
42
+ device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
43
+ }
44
+
45
+ std::string toString() const {
46
+ std::stringstream str;
47
+ str << "IVFFlat device " << device << " numVecs " << numAdd << " dim "
48
+ << dim << " numCentroids " << numCentroids << " nprobe " << nprobe
49
+ << " numQuery " << numQuery << " k " << k << " indicesOpt "
50
+ << indicesOpt;
51
+
52
+ return str.str();
53
+ }
54
+
55
+ int numAdd;
56
+ int dim;
57
+ int numCentroids;
58
+ int numTrain;
59
+ int nprobe;
60
+ int numQuery;
61
+ int k;
62
+ int device;
63
+ faiss::gpu::IndicesOptions indicesOpt;
70
64
  };
71
65
 
72
- void queryTest(faiss::MetricType metricType,
73
- bool useFloat16CoarseQuantizer,
74
- int dimOverride = -1) {
75
- for (int tries = 0; tries < 2; ++tries) {
76
- Options opt;
77
- opt.dim = dimOverride != -1 ? dimOverride : opt.dim;
66
+ void queryTest(
67
+ faiss::MetricType metricType,
68
+ bool useFloat16CoarseQuantizer,
69
+ int dimOverride = -1) {
70
+ for (int tries = 0; tries < 2; ++tries) {
71
+ Options opt;
72
+ opt.dim = dimOverride != -1 ? dimOverride : opt.dim;
73
+
74
+ std::vector<float> trainVecs =
75
+ faiss::gpu::randVecs(opt.numTrain, opt.dim);
76
+ std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
77
+
78
+ faiss::IndexFlatL2 quantizerL2(opt.dim);
79
+ faiss::IndexFlatIP quantizerIP(opt.dim);
80
+ faiss::Index* quantizer = metricType == faiss::METRIC_L2
81
+ ? (faiss::Index*)&quantizerL2
82
+ : (faiss::Index*)&quantizerIP;
83
+
84
+ faiss::IndexIVFFlat cpuIndex(
85
+ quantizer, opt.dim, opt.numCentroids, metricType);
86
+ cpuIndex.train(opt.numTrain, trainVecs.data());
87
+ cpuIndex.add(opt.numAdd, addVecs.data());
88
+ cpuIndex.nprobe = opt.nprobe;
89
+
90
+ faiss::gpu::StandardGpuResources res;
91
+ res.noTempMemory();
92
+
93
+ faiss::gpu::GpuIndexIVFFlatConfig config;
94
+ config.device = opt.device;
95
+ config.indicesOptions = opt.indicesOpt;
96
+ config.flatConfig.useFloat16 = useFloat16CoarseQuantizer;
97
+
98
+ faiss::gpu::GpuIndexIVFFlat gpuIndex(
99
+ &res, cpuIndex.d, cpuIndex.nlist, cpuIndex.metric_type, config);
100
+ gpuIndex.copyFrom(&cpuIndex);
101
+ gpuIndex.setNumProbes(opt.nprobe);
102
+
103
+ bool compFloat16 = useFloat16CoarseQuantizer;
104
+ faiss::gpu::compareIndices(
105
+ cpuIndex,
106
+ gpuIndex,
107
+ opt.numQuery,
108
+ opt.dim,
109
+ opt.k,
110
+ opt.toString(),
111
+ compFloat16 ? kF16MaxRelErr : kF32MaxRelErr,
112
+ // FIXME: the fp16 bounds are
113
+ // useless when math (the accumulator) is
114
+ // in fp16. Figure out another way to test
115
+ compFloat16 ? 0.70f : 0.1f,
116
+ compFloat16 ? 0.65f : 0.015f);
117
+ }
118
+ }
78
119
 
120
+ void addTest(faiss::MetricType metricType, bool useFloat16CoarseQuantizer) {
121
+ for (int tries = 0; tries < 2; ++tries) {
122
+ Options opt;
123
+
124
+ std::vector<float> trainVecs =
125
+ faiss::gpu::randVecs(opt.numTrain, opt.dim);
126
+ std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
127
+
128
+ faiss::IndexFlatL2 quantizerL2(opt.dim);
129
+ faiss::IndexFlatIP quantizerIP(opt.dim);
130
+ faiss::Index* quantizer = metricType == faiss::METRIC_L2
131
+ ? (faiss::Index*)&quantizerL2
132
+ : (faiss::Index*)&quantizerIP;
133
+
134
+ faiss::IndexIVFFlat cpuIndex(
135
+ quantizer, opt.dim, opt.numCentroids, metricType);
136
+ cpuIndex.train(opt.numTrain, trainVecs.data());
137
+ cpuIndex.nprobe = opt.nprobe;
138
+
139
+ faiss::gpu::StandardGpuResources res;
140
+ res.noTempMemory();
141
+
142
+ faiss::gpu::GpuIndexIVFFlatConfig config;
143
+ config.device = opt.device;
144
+ config.indicesOptions = opt.indicesOpt;
145
+ config.flatConfig.useFloat16 = useFloat16CoarseQuantizer;
146
+
147
+ faiss::gpu::GpuIndexIVFFlat gpuIndex(
148
+ &res, cpuIndex.d, cpuIndex.nlist, cpuIndex.metric_type, config);
149
+ gpuIndex.copyFrom(&cpuIndex);
150
+ gpuIndex.setNumProbes(opt.nprobe);
151
+
152
+ cpuIndex.add(opt.numAdd, addVecs.data());
153
+ gpuIndex.add(opt.numAdd, addVecs.data());
154
+
155
+ bool compFloat16 = useFloat16CoarseQuantizer;
156
+ faiss::gpu::compareIndices(
157
+ cpuIndex,
158
+ gpuIndex,
159
+ opt.numQuery,
160
+ opt.dim,
161
+ opt.k,
162
+ opt.toString(),
163
+ compFloat16 ? kF16MaxRelErr : kF32MaxRelErr,
164
+ compFloat16 ? 0.70f : 0.1f,
165
+ compFloat16 ? 0.30f : 0.015f);
166
+ }
167
+ }
168
+
169
+ void copyToTest(bool useFloat16CoarseQuantizer) {
170
+ Options opt;
79
171
  std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
80
172
  std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
81
173
 
82
- faiss::IndexFlatL2 quantizerL2(opt.dim);
83
- faiss::IndexFlatIP quantizerIP(opt.dim);
84
- faiss::Index* quantizer =
85
- metricType == faiss::METRIC_L2 ?
86
- (faiss::Index*) &quantizerL2 : (faiss::Index*) &quantizerIP;
87
-
88
- faiss::IndexIVFFlat cpuIndex(quantizer,
89
- opt.dim, opt.numCentroids, metricType);
90
- cpuIndex.train(opt.numTrain, trainVecs.data());
91
- cpuIndex.add(opt.numAdd, addVecs.data());
92
- cpuIndex.nprobe = opt.nprobe;
93
-
94
174
  faiss::gpu::StandardGpuResources res;
95
175
  res.noTempMemory();
96
176
 
@@ -99,47 +179,57 @@ void queryTest(faiss::MetricType metricType,
99
179
  config.indicesOptions = opt.indicesOpt;
100
180
  config.flatConfig.useFloat16 = useFloat16CoarseQuantizer;
101
181
 
102
- faiss::gpu::GpuIndexIVFFlat gpuIndex(&res,
103
- cpuIndex.d,
104
- cpuIndex.nlist,
105
- cpuIndex.metric_type,
106
- config);
107
- gpuIndex.copyFrom(&cpuIndex);
182
+ faiss::gpu::GpuIndexIVFFlat gpuIndex(
183
+ &res, opt.dim, opt.numCentroids, faiss::METRIC_L2, config);
184
+ gpuIndex.train(opt.numTrain, trainVecs.data());
185
+ gpuIndex.add(opt.numAdd, addVecs.data());
108
186
  gpuIndex.setNumProbes(opt.nprobe);
109
187
 
188
+ // use garbage values to see if we overwrite then
189
+ faiss::IndexFlatL2 cpuQuantizer(1);
190
+ faiss::IndexIVFFlat cpuIndex(&cpuQuantizer, 1, 1, faiss::METRIC_L2);
191
+ cpuIndex.nprobe = 1;
192
+
193
+ gpuIndex.copyTo(&cpuIndex);
194
+
195
+ EXPECT_EQ(cpuIndex.ntotal, gpuIndex.ntotal);
196
+ EXPECT_EQ(gpuIndex.ntotal, opt.numAdd);
197
+
198
+ EXPECT_EQ(cpuIndex.d, gpuIndex.d);
199
+ EXPECT_EQ(cpuIndex.quantizer->d, gpuIndex.quantizer->d);
200
+ EXPECT_EQ(cpuIndex.d, opt.dim);
201
+ EXPECT_EQ(cpuIndex.nlist, gpuIndex.getNumLists());
202
+ EXPECT_EQ(cpuIndex.nprobe, gpuIndex.getNumProbes());
203
+
204
+ testIVFEquality(cpuIndex, gpuIndex);
205
+
206
+ // Query both objects; results should be equivalent
110
207
  bool compFloat16 = useFloat16CoarseQuantizer;
111
- faiss::gpu::compareIndices(cpuIndex, gpuIndex,
112
- opt.numQuery, opt.dim, opt.k, opt.toString(),
113
- compFloat16 ? kF16MaxRelErr : kF32MaxRelErr,
114
- // FIXME: the fp16 bounds are
115
- // useless when math (the accumulator) is
116
- // in fp16. Figure out another way to test
117
- compFloat16 ? 0.70f : 0.1f,
118
- compFloat16 ? 0.65f : 0.015f);
119
- }
208
+ faiss::gpu::compareIndices(
209
+ cpuIndex,
210
+ gpuIndex,
211
+ opt.numQuery,
212
+ opt.dim,
213
+ opt.k,
214
+ opt.toString(),
215
+ compFloat16 ? kF16MaxRelErr : kF32MaxRelErr,
216
+ compFloat16 ? 0.70f : 0.1f,
217
+ compFloat16 ? 0.30f : 0.015f);
120
218
  }
121
219
 
122
- void addTest(faiss::MetricType metricType,
123
- bool useFloat16CoarseQuantizer) {
124
- for (int tries = 0; tries < 2; ++tries) {
220
+ void copyFromTest(bool useFloat16CoarseQuantizer) {
125
221
  Options opt;
126
-
127
222
  std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
128
223
  std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
129
224
 
130
- faiss::IndexFlatL2 quantizerL2(opt.dim);
131
- faiss::IndexFlatIP quantizerIP(opt.dim);
132
- faiss::Index* quantizer =
133
- metricType == faiss::METRIC_L2 ?
134
- (faiss::Index*) &quantizerL2 : (faiss::Index*) &quantizerIP;
135
-
136
- faiss::IndexIVFFlat cpuIndex(quantizer,
137
- opt.dim,
138
- opt.numCentroids,
139
- metricType);
140
- cpuIndex.train(opt.numTrain, trainVecs.data());
225
+ faiss::IndexFlatL2 cpuQuantizer(opt.dim);
226
+ faiss::IndexIVFFlat cpuIndex(
227
+ &cpuQuantizer, opt.dim, opt.numCentroids, faiss::METRIC_L2);
141
228
  cpuIndex.nprobe = opt.nprobe;
229
+ cpuIndex.train(opt.numTrain, trainVecs.data());
230
+ cpuIndex.add(opt.numAdd, addVecs.data());
142
231
 
232
+ // use garbage values to see if we overwrite then
143
233
  faiss::gpu::StandardGpuResources res;
144
234
  res.noTempMemory();
145
235
 
@@ -148,140 +238,49 @@ void addTest(faiss::MetricType metricType,
148
238
  config.indicesOptions = opt.indicesOpt;
149
239
  config.flatConfig.useFloat16 = useFloat16CoarseQuantizer;
150
240
 
151
- faiss::gpu::GpuIndexIVFFlat gpuIndex(&res,
152
- cpuIndex.d,
153
- cpuIndex.nlist,
154
- cpuIndex.metric_type,
155
- config);
241
+ faiss::gpu::GpuIndexIVFFlat gpuIndex(&res, 1, 1, faiss::METRIC_L2, config);
242
+ gpuIndex.setNumProbes(1);
243
+
156
244
  gpuIndex.copyFrom(&cpuIndex);
157
- gpuIndex.setNumProbes(opt.nprobe);
158
245
 
159
- cpuIndex.add(opt.numAdd, addVecs.data());
160
- gpuIndex.add(opt.numAdd, addVecs.data());
246
+ EXPECT_EQ(cpuIndex.ntotal, gpuIndex.ntotal);
247
+ EXPECT_EQ(gpuIndex.ntotal, opt.numAdd);
161
248
 
162
- bool compFloat16 = useFloat16CoarseQuantizer;
163
- faiss::gpu::compareIndices(cpuIndex, gpuIndex,
164
- opt.numQuery, opt.dim, opt.k, opt.toString(),
165
- compFloat16 ? kF16MaxRelErr : kF32MaxRelErr,
166
- compFloat16 ? 0.70f : 0.1f,
167
- compFloat16 ? 0.30f : 0.015f);
168
- }
169
- }
249
+ EXPECT_EQ(cpuIndex.d, gpuIndex.d);
250
+ EXPECT_EQ(cpuIndex.d, opt.dim);
251
+ EXPECT_EQ(cpuIndex.nlist, gpuIndex.getNumLists());
252
+ EXPECT_EQ(cpuIndex.nprobe, gpuIndex.getNumProbes());
170
253
 
171
- void copyToTest(bool useFloat16CoarseQuantizer) {
172
- Options opt;
173
- std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
174
- std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
175
-
176
- faiss::gpu::StandardGpuResources res;
177
- res.noTempMemory();
178
-
179
- faiss::gpu::GpuIndexIVFFlatConfig config;
180
- config.device = opt.device;
181
- config.indicesOptions = opt.indicesOpt;
182
- config.flatConfig.useFloat16 = useFloat16CoarseQuantizer;
183
-
184
- faiss::gpu::GpuIndexIVFFlat gpuIndex(&res,
185
- opt.dim,
186
- opt.numCentroids,
187
- faiss::METRIC_L2,
188
- config);
189
- gpuIndex.train(opt.numTrain, trainVecs.data());
190
- gpuIndex.add(opt.numAdd, addVecs.data());
191
- gpuIndex.setNumProbes(opt.nprobe);
192
-
193
- // use garbage values to see if we overwrite then
194
- faiss::IndexFlatL2 cpuQuantizer(1);
195
- faiss::IndexIVFFlat cpuIndex(&cpuQuantizer, 1, 1, faiss::METRIC_L2);
196
- cpuIndex.nprobe = 1;
197
-
198
- gpuIndex.copyTo(&cpuIndex);
199
-
200
- EXPECT_EQ(cpuIndex.ntotal, gpuIndex.ntotal);
201
- EXPECT_EQ(gpuIndex.ntotal, opt.numAdd);
202
-
203
- EXPECT_EQ(cpuIndex.d, gpuIndex.d);
204
- EXPECT_EQ(cpuIndex.quantizer->d, gpuIndex.quantizer->d);
205
- EXPECT_EQ(cpuIndex.d, opt.dim);
206
- EXPECT_EQ(cpuIndex.nlist, gpuIndex.getNumLists());
207
- EXPECT_EQ(cpuIndex.nprobe, gpuIndex.getNumProbes());
208
-
209
- testIVFEquality(cpuIndex, gpuIndex);
210
-
211
- // Query both objects; results should be equivalent
212
- bool compFloat16 = useFloat16CoarseQuantizer;
213
- faiss::gpu::compareIndices(cpuIndex, gpuIndex,
214
- opt.numQuery, opt.dim, opt.k, opt.toString(),
215
- compFloat16 ? kF16MaxRelErr : kF32MaxRelErr,
216
- compFloat16 ? 0.70f : 0.1f,
217
- compFloat16 ? 0.30f : 0.015f);
218
- }
254
+ testIVFEquality(cpuIndex, gpuIndex);
219
255
 
220
- void copyFromTest(bool useFloat16CoarseQuantizer) {
221
- Options opt;
222
- std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
223
- std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
224
-
225
- faiss::IndexFlatL2 cpuQuantizer(opt.dim);
226
- faiss::IndexIVFFlat cpuIndex(&cpuQuantizer,
227
- opt.dim,
228
- opt.numCentroids,
229
- faiss::METRIC_L2);
230
- cpuIndex.nprobe = opt.nprobe;
231
- cpuIndex.train(opt.numTrain, trainVecs.data());
232
- cpuIndex.add(opt.numAdd, addVecs.data());
233
-
234
- // use garbage values to see if we overwrite then
235
- faiss::gpu::StandardGpuResources res;
236
- res.noTempMemory();
237
-
238
- faiss::gpu::GpuIndexIVFFlatConfig config;
239
- config.device = opt.device;
240
- config.indicesOptions = opt.indicesOpt;
241
- config.flatConfig.useFloat16 = useFloat16CoarseQuantizer;
242
-
243
- faiss::gpu::GpuIndexIVFFlat gpuIndex(&res,
244
- 1,
245
- 1,
246
- faiss::METRIC_L2,
247
- config);
248
- gpuIndex.setNumProbes(1);
249
-
250
- gpuIndex.copyFrom(&cpuIndex);
251
-
252
- EXPECT_EQ(cpuIndex.ntotal, gpuIndex.ntotal);
253
- EXPECT_EQ(gpuIndex.ntotal, opt.numAdd);
254
-
255
- EXPECT_EQ(cpuIndex.d, gpuIndex.d);
256
- EXPECT_EQ(cpuIndex.d, opt.dim);
257
- EXPECT_EQ(cpuIndex.nlist, gpuIndex.getNumLists());
258
- EXPECT_EQ(cpuIndex.nprobe, gpuIndex.getNumProbes());
259
-
260
- testIVFEquality(cpuIndex, gpuIndex);
261
-
262
- // Query both objects; results should be equivalent
263
- bool compFloat16 = useFloat16CoarseQuantizer;
264
- faiss::gpu::compareIndices(cpuIndex, gpuIndex,
265
- opt.numQuery, opt.dim, opt.k, opt.toString(),
266
- compFloat16 ? kF16MaxRelErr : kF32MaxRelErr,
267
- compFloat16 ? 0.70f : 0.1f,
268
- compFloat16 ? 0.30f : 0.015f);
256
+ // Query both objects; results should be equivalent
257
+ bool compFloat16 = useFloat16CoarseQuantizer;
258
+ faiss::gpu::compareIndices(
259
+ cpuIndex,
260
+ gpuIndex,
261
+ opt.numQuery,
262
+ opt.dim,
263
+ opt.k,
264
+ opt.toString(),
265
+ compFloat16 ? kF16MaxRelErr : kF32MaxRelErr,
266
+ compFloat16 ? 0.70f : 0.1f,
267
+ compFloat16 ? 0.30f : 0.015f);
269
268
  }
270
269
 
271
270
  TEST(TestGpuIndexIVFFlat, Float32_32_Add_L2) {
272
- addTest(faiss::METRIC_L2, false);
271
+ addTest(faiss::METRIC_L2, false);
273
272
  }
274
273
 
275
274
  TEST(TestGpuIndexIVFFlat, Float32_32_Add_IP) {
276
- addTest(faiss::METRIC_INNER_PRODUCT, false);
275
+ addTest(faiss::METRIC_INNER_PRODUCT, false);
277
276
  }
278
277
 
279
278
  TEST(TestGpuIndexIVFFlat, Float16_32_Add_L2) {
280
- addTest(faiss::METRIC_L2, true);
279
+ addTest(faiss::METRIC_L2, true);
281
280
  }
282
281
 
283
282
  TEST(TestGpuIndexIVFFlat, Float16_32_Add_IP) {
284
- addTest(faiss::METRIC_INNER_PRODUCT, true);
283
+ addTest(faiss::METRIC_INNER_PRODUCT, true);
285
284
  }
286
285
 
287
286
  //
@@ -289,21 +288,21 @@ TEST(TestGpuIndexIVFFlat, Float16_32_Add_IP) {
289
288
  //
290
289
 
291
290
  TEST(TestGpuIndexIVFFlat, Float32_Query_L2) {
292
- queryTest(faiss::METRIC_L2, false);
291
+ queryTest(faiss::METRIC_L2, false);
293
292
  }
294
293
 
295
294
  TEST(TestGpuIndexIVFFlat, Float32_Query_IP) {
296
- queryTest(faiss::METRIC_INNER_PRODUCT, false);
295
+ queryTest(faiss::METRIC_INNER_PRODUCT, false);
297
296
  }
298
297
 
299
298
  // float16 coarse quantizer
300
299
 
301
300
  TEST(TestGpuIndexIVFFlat, Float16_32_Query_L2) {
302
- queryTest(faiss::METRIC_L2, true);
301
+ queryTest(faiss::METRIC_L2, true);
303
302
  }
304
303
 
305
304
  TEST(TestGpuIndexIVFFlat, Float16_32_Query_IP) {
306
- queryTest(faiss::METRIC_INNER_PRODUCT, true);
305
+ queryTest(faiss::METRIC_INNER_PRODUCT, true);
307
306
  }
308
307
 
309
308
  //
@@ -312,19 +311,19 @@ TEST(TestGpuIndexIVFFlat, Float16_32_Query_IP) {
312
311
  //
313
312
 
314
313
  TEST(TestGpuIndexIVFFlat, Float32_Query_L2_64) {
315
- queryTest(faiss::METRIC_L2, false, 64);
314
+ queryTest(faiss::METRIC_L2, false, 64);
316
315
  }
317
316
 
318
317
  TEST(TestGpuIndexIVFFlat, Float32_Query_IP_64) {
319
- queryTest(faiss::METRIC_INNER_PRODUCT, false, 64);
318
+ queryTest(faiss::METRIC_INNER_PRODUCT, false, 64);
320
319
  }
321
320
 
322
321
  TEST(TestGpuIndexIVFFlat, Float32_Query_L2_128) {
323
- queryTest(faiss::METRIC_L2, false, 128);
322
+ queryTest(faiss::METRIC_L2, false, 128);
324
323
  }
325
324
 
326
325
  TEST(TestGpuIndexIVFFlat, Float32_Query_IP_128) {
327
- queryTest(faiss::METRIC_INNER_PRODUCT, false, 128);
326
+ queryTest(faiss::METRIC_INNER_PRODUCT, false, 128);
328
327
  }
329
328
 
330
329
  //
@@ -332,71 +331,72 @@ TEST(TestGpuIndexIVFFlat, Float32_Query_IP_128) {
332
331
  //
333
332
 
334
333
  TEST(TestGpuIndexIVFFlat, Float32_32_CopyTo) {
335
- copyToTest(false);
334
+ copyToTest(false);
336
335
  }
337
336
 
338
337
  TEST(TestGpuIndexIVFFlat, Float32_32_CopyFrom) {
339
- copyFromTest(false);
338
+ copyFromTest(false);
340
339
  }
341
340
 
342
341
  TEST(TestGpuIndexIVFFlat, Float32_negative) {
343
- Options opt;
344
-
345
- auto trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
346
- auto addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
347
-
348
- // Put all vecs on negative side
349
- for (auto& f : trainVecs) {
350
- f = std::abs(f) * -1.0f;
351
- }
352
-
353
- for (auto& f : addVecs) {
354
- f *= std::abs(f) * -1.0f;
355
- }
356
-
357
- faiss::IndexFlatIP quantizerIP(opt.dim);
358
- faiss::Index* quantizer = (faiss::Index*) &quantizerIP;
359
-
360
- faiss::IndexIVFFlat cpuIndex(quantizer,
361
- opt.dim, opt.numCentroids,
362
- faiss::METRIC_INNER_PRODUCT);
363
- cpuIndex.train(opt.numTrain, trainVecs.data());
364
- cpuIndex.add(opt.numAdd, addVecs.data());
365
- cpuIndex.nprobe = opt.nprobe;
366
-
367
- faiss::gpu::StandardGpuResources res;
368
- res.noTempMemory();
369
-
370
- faiss::gpu::GpuIndexIVFFlatConfig config;
371
- config.device = opt.device;
372
- config.indicesOptions = opt.indicesOpt;
373
-
374
- faiss::gpu::GpuIndexIVFFlat gpuIndex(&res,
375
- cpuIndex.d,
376
- cpuIndex.nlist,
377
- cpuIndex.metric_type,
378
- config);
379
- gpuIndex.copyFrom(&cpuIndex);
380
- gpuIndex.setNumProbes(opt.nprobe);
381
-
382
- // Construct a positive test set
383
- auto queryVecs = faiss::gpu::randVecs(opt.numQuery, opt.dim);
384
-
385
- // Put all vecs on positive size
386
- for (auto& f : queryVecs) {
387
- f = std::abs(f);
388
- }
389
-
390
- bool compFloat16 = false;
391
- faiss::gpu::compareIndices(queryVecs,
392
- cpuIndex, gpuIndex,
393
- opt.numQuery, opt.dim, opt.k, opt.toString(),
394
- compFloat16 ? kF16MaxRelErr : kF32MaxRelErr,
395
- // FIXME: the fp16 bounds are
396
- // useless when math (the accumulator) is
397
- // in fp16. Figure out another way to test
398
- compFloat16 ? 0.99f : 0.1f,
399
- compFloat16 ? 0.65f : 0.015f);
342
+ Options opt;
343
+
344
+ auto trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
345
+ auto addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
346
+
347
+ // Put all vecs on negative side
348
+ for (auto& f : trainVecs) {
349
+ f = std::abs(f) * -1.0f;
350
+ }
351
+
352
+ for (auto& f : addVecs) {
353
+ f *= std::abs(f) * -1.0f;
354
+ }
355
+
356
+ faiss::IndexFlatIP quantizerIP(opt.dim);
357
+ faiss::Index* quantizer = (faiss::Index*)&quantizerIP;
358
+
359
+ faiss::IndexIVFFlat cpuIndex(
360
+ quantizer, opt.dim, opt.numCentroids, faiss::METRIC_INNER_PRODUCT);
361
+ cpuIndex.train(opt.numTrain, trainVecs.data());
362
+ cpuIndex.add(opt.numAdd, addVecs.data());
363
+ cpuIndex.nprobe = opt.nprobe;
364
+
365
+ faiss::gpu::StandardGpuResources res;
366
+ res.noTempMemory();
367
+
368
+ faiss::gpu::GpuIndexIVFFlatConfig config;
369
+ config.device = opt.device;
370
+ config.indicesOptions = opt.indicesOpt;
371
+
372
+ faiss::gpu::GpuIndexIVFFlat gpuIndex(
373
+ &res, cpuIndex.d, cpuIndex.nlist, cpuIndex.metric_type, config);
374
+ gpuIndex.copyFrom(&cpuIndex);
375
+ gpuIndex.setNumProbes(opt.nprobe);
376
+
377
+ // Construct a positive test set
378
+ auto queryVecs = faiss::gpu::randVecs(opt.numQuery, opt.dim);
379
+
380
+ // Put all vecs on positive size
381
+ for (auto& f : queryVecs) {
382
+ f = std::abs(f);
383
+ }
384
+
385
+ bool compFloat16 = false;
386
+ faiss::gpu::compareIndices(
387
+ queryVecs,
388
+ cpuIndex,
389
+ gpuIndex,
390
+ opt.numQuery,
391
+ opt.dim,
392
+ opt.k,
393
+ opt.toString(),
394
+ compFloat16 ? kF16MaxRelErr : kF32MaxRelErr,
395
+ // FIXME: the fp16 bounds are
396
+ // useless when math (the accumulator) is
397
+ // in fp16. Figure out another way to test
398
+ compFloat16 ? 0.99f : 0.1f,
399
+ compFloat16 ? 0.65f : 0.015f);
400
400
  }
401
401
 
402
402
  //
@@ -404,152 +404,152 @@ TEST(TestGpuIndexIVFFlat, Float32_negative) {
404
404
  //
405
405
 
406
406
  TEST(TestGpuIndexIVFFlat, QueryNaN) {
407
- Options opt;
408
-
409
- std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
410
- std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
411
-
412
- faiss::gpu::StandardGpuResources res;
413
- res.noTempMemory();
414
-
415
- faiss::gpu::GpuIndexIVFFlatConfig config;
416
- config.device = opt.device;
417
- config.indicesOptions = opt.indicesOpt;
418
- config.flatConfig.useFloat16 = faiss::gpu::randBool();
419
-
420
- faiss::gpu::GpuIndexIVFFlat gpuIndex(&res,
421
- opt.dim,
422
- opt.numCentroids,
423
- faiss::METRIC_L2,
424
- config);
425
- gpuIndex.setNumProbes(opt.nprobe);
426
-
427
- gpuIndex.train(opt.numTrain, trainVecs.data());
428
- gpuIndex.add(opt.numAdd, addVecs.data());
429
-
430
- int numQuery = 10;
431
- std::vector<float> nans(numQuery * opt.dim,
432
- std::numeric_limits<float>::quiet_NaN());
433
-
434
- std::vector<float> distances(numQuery * opt.k, 0);
435
- std::vector<faiss::Index::idx_t> indices(numQuery * opt.k, 0);
436
-
437
- gpuIndex.search(numQuery,
438
- nans.data(),
439
- opt.k,
440
- distances.data(),
441
- indices.data());
442
-
443
- for (int q = 0; q < numQuery; ++q) {
444
- for (int k = 0; k < opt.k; ++k) {
445
- EXPECT_EQ(indices[q * opt.k + k], -1);
446
- EXPECT_EQ(distances[q * opt.k + k], std::numeric_limits<float>::max());
407
+ Options opt;
408
+
409
+ std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
410
+ std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
411
+
412
+ faiss::gpu::StandardGpuResources res;
413
+ res.noTempMemory();
414
+
415
+ faiss::gpu::GpuIndexIVFFlatConfig config;
416
+ config.device = opt.device;
417
+ config.indicesOptions = opt.indicesOpt;
418
+ config.flatConfig.useFloat16 = faiss::gpu::randBool();
419
+
420
+ faiss::gpu::GpuIndexIVFFlat gpuIndex(
421
+ &res, opt.dim, opt.numCentroids, faiss::METRIC_L2, config);
422
+ gpuIndex.setNumProbes(opt.nprobe);
423
+
424
+ gpuIndex.train(opt.numTrain, trainVecs.data());
425
+ gpuIndex.add(opt.numAdd, addVecs.data());
426
+
427
+ int numQuery = 10;
428
+ std::vector<float> nans(
429
+ numQuery * opt.dim, std::numeric_limits<float>::quiet_NaN());
430
+
431
+ std::vector<float> distances(numQuery * opt.k, 0);
432
+ std::vector<faiss::Index::idx_t> indices(numQuery * opt.k, 0);
433
+
434
+ gpuIndex.search(
435
+ numQuery, nans.data(), opt.k, distances.data(), indices.data());
436
+
437
+ for (int q = 0; q < numQuery; ++q) {
438
+ for (int k = 0; k < opt.k; ++k) {
439
+ EXPECT_EQ(indices[q * opt.k + k], -1);
440
+ EXPECT_EQ(
441
+ distances[q * opt.k + k],
442
+ std::numeric_limits<float>::max());
443
+ }
447
444
  }
448
- }
449
445
  }
450
446
 
451
447
  TEST(TestGpuIndexIVFFlat, AddNaN) {
452
- Options opt;
453
-
454
- faiss::gpu::StandardGpuResources res;
455
- res.noTempMemory();
456
-
457
- faiss::gpu::GpuIndexIVFFlatConfig config;
458
- config.device = opt.device;
459
- config.indicesOptions = opt.indicesOpt;
460
- config.flatConfig.useFloat16 = faiss::gpu::randBool();
461
-
462
- faiss::gpu::GpuIndexIVFFlat gpuIndex(&res,
463
- opt.dim,
464
- opt.numCentroids,
465
- faiss::METRIC_L2,
466
- config);
467
- gpuIndex.setNumProbes(opt.nprobe);
468
-
469
- int numNans = 10;
470
- std::vector<float> nans(numNans * opt.dim,
471
- std::numeric_limits<float>::quiet_NaN());
472
-
473
- // Make one vector valid (not the first vector, in order to test offset
474
- // issues), which should actually add
475
- for (int i = 0; i < opt.dim; ++i) {
476
- nans[opt.dim + i] = i;
477
- }
478
-
479
- std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
480
- gpuIndex.train(opt.numTrain, trainVecs.data());
481
-
482
- // should not crash
483
- EXPECT_EQ(gpuIndex.ntotal, 0);
484
- gpuIndex.add(numNans, nans.data());
485
-
486
- std::vector<float> queryVecs = faiss::gpu::randVecs(opt.numQuery, opt.dim);
487
- std::vector<float> distance(opt.numQuery * opt.k, 0);
488
- std::vector<faiss::Index::idx_t> indices(opt.numQuery * opt.k, 0);
489
-
490
- // should not crash
491
- gpuIndex.search(opt.numQuery, queryVecs.data(), opt.k,
492
- distance.data(), indices.data());
448
+ Options opt;
449
+
450
+ faiss::gpu::StandardGpuResources res;
451
+ res.noTempMemory();
452
+
453
+ faiss::gpu::GpuIndexIVFFlatConfig config;
454
+ config.device = opt.device;
455
+ config.indicesOptions = opt.indicesOpt;
456
+ config.flatConfig.useFloat16 = faiss::gpu::randBool();
457
+
458
+ faiss::gpu::GpuIndexIVFFlat gpuIndex(
459
+ &res, opt.dim, opt.numCentroids, faiss::METRIC_L2, config);
460
+ gpuIndex.setNumProbes(opt.nprobe);
461
+
462
+ int numNans = 10;
463
+ std::vector<float> nans(
464
+ numNans * opt.dim, std::numeric_limits<float>::quiet_NaN());
465
+
466
+ // Make one vector valid (not the first vector, in order to test offset
467
+ // issues), which should actually add
468
+ for (int i = 0; i < opt.dim; ++i) {
469
+ nans[opt.dim + i] = i;
470
+ }
471
+
472
+ std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
473
+ gpuIndex.train(opt.numTrain, trainVecs.data());
474
+
475
+ // should not crash
476
+ EXPECT_EQ(gpuIndex.ntotal, 0);
477
+ gpuIndex.add(numNans, nans.data());
478
+
479
+ std::vector<float> queryVecs = faiss::gpu::randVecs(opt.numQuery, opt.dim);
480
+ std::vector<float> distance(opt.numQuery * opt.k, 0);
481
+ std::vector<faiss::Index::idx_t> indices(opt.numQuery * opt.k, 0);
482
+
483
+ // should not crash
484
+ gpuIndex.search(
485
+ opt.numQuery,
486
+ queryVecs.data(),
487
+ opt.k,
488
+ distance.data(),
489
+ indices.data());
493
490
  }
494
491
 
495
492
  TEST(TestGpuIndexIVFFlat, UnifiedMemory) {
496
- // Construct on a random device to test multi-device, if we have
497
- // multiple devices
498
- int device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
499
-
500
- if (!faiss::gpu::getFullUnifiedMemSupport(device)) {
501
- return;
502
- }
503
-
504
- int dim = 128;
505
-
506
- int numCentroids = 256;
507
- // Unfortunately it would take forever to add 24 GB in IVFPQ data,
508
- // so just perform a small test with data allocated in the unified
509
- // memory address space
510
- size_t numAdd = 10000;
511
- size_t numTrain = numCentroids * 40;
512
- int numQuery = 10;
513
- int k = 10;
514
- int nprobe = 8;
515
-
516
- std::vector<float> trainVecs = faiss::gpu::randVecs(numTrain, dim);
517
- std::vector<float> addVecs = faiss::gpu::randVecs(numAdd, dim);
518
-
519
- faiss::IndexFlatL2 quantizer(dim);
520
- faiss::IndexIVFFlat cpuIndex(&quantizer, dim, numCentroids, faiss::METRIC_L2);
521
-
522
- cpuIndex.train(numTrain, trainVecs.data());
523
- cpuIndex.add(numAdd, addVecs.data());
524
- cpuIndex.nprobe = nprobe;
525
-
526
- faiss::gpu::StandardGpuResources res;
527
- res.noTempMemory();
528
-
529
- faiss::gpu::GpuIndexIVFFlatConfig config;
530
- config.device = device;
531
- config.memorySpace = faiss::gpu::MemorySpace::Unified;
532
-
533
- faiss::gpu::GpuIndexIVFFlat gpuIndex(&res,
534
- dim,
535
- numCentroids,
536
- faiss::METRIC_L2,
537
- config);
538
- gpuIndex.copyFrom(&cpuIndex);
539
- gpuIndex.setNumProbes(nprobe);
540
-
541
- faiss::gpu::compareIndices(cpuIndex, gpuIndex,
542
- numQuery, dim, k, "Unified Memory",
543
- kF32MaxRelErr,
544
- 0.1f,
545
- 0.015f);
493
+ // Construct on a random device to test multi-device, if we have
494
+ // multiple devices
495
+ int device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
496
+
497
+ if (!faiss::gpu::getFullUnifiedMemSupport(device)) {
498
+ return;
499
+ }
500
+
501
+ int dim = 128;
502
+
503
+ int numCentroids = 256;
504
+ // Unfortunately it would take forever to add 24 GB in IVFPQ data,
505
+ // so just perform a small test with data allocated in the unified
506
+ // memory address space
507
+ size_t numAdd = 10000;
508
+ size_t numTrain = numCentroids * 40;
509
+ int numQuery = 10;
510
+ int k = 10;
511
+ int nprobe = 8;
512
+
513
+ std::vector<float> trainVecs = faiss::gpu::randVecs(numTrain, dim);
514
+ std::vector<float> addVecs = faiss::gpu::randVecs(numAdd, dim);
515
+
516
+ faiss::IndexFlatL2 quantizer(dim);
517
+ faiss::IndexIVFFlat cpuIndex(
518
+ &quantizer, dim, numCentroids, faiss::METRIC_L2);
519
+
520
+ cpuIndex.train(numTrain, trainVecs.data());
521
+ cpuIndex.add(numAdd, addVecs.data());
522
+ cpuIndex.nprobe = nprobe;
523
+
524
+ faiss::gpu::StandardGpuResources res;
525
+ res.noTempMemory();
526
+
527
+ faiss::gpu::GpuIndexIVFFlatConfig config;
528
+ config.device = device;
529
+ config.memorySpace = faiss::gpu::MemorySpace::Unified;
530
+
531
+ faiss::gpu::GpuIndexIVFFlat gpuIndex(
532
+ &res, dim, numCentroids, faiss::METRIC_L2, config);
533
+ gpuIndex.copyFrom(&cpuIndex);
534
+ gpuIndex.setNumProbes(nprobe);
535
+
536
+ faiss::gpu::compareIndices(
537
+ cpuIndex,
538
+ gpuIndex,
539
+ numQuery,
540
+ dim,
541
+ k,
542
+ "Unified Memory",
543
+ kF32MaxRelErr,
544
+ 0.1f,
545
+ 0.015f);
546
546
  }
547
547
 
548
548
  int main(int argc, char** argv) {
549
- testing::InitGoogleTest(&argc, argv);
549
+ testing::InitGoogleTest(&argc, argv);
550
550
 
551
- // just run with a fixed test seed
552
- faiss::gpu::setTestSeed(100);
551
+ // just run with a fixed test seed
552
+ faiss::gpu::setTestSeed(100);
553
553
 
554
- return RUN_ALL_TESTS();
554
+ return RUN_ALL_TESTS();
555
555
  }