faiss 0.2.0 → 0.2.4

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (215) hide show
  1. checksums.yaml +4 -4
  2. data/CHANGELOG.md +16 -0
  3. data/LICENSE.txt +1 -1
  4. data/README.md +7 -7
  5. data/ext/faiss/extconf.rb +6 -3
  6. data/ext/faiss/numo.hpp +4 -4
  7. data/ext/faiss/utils.cpp +1 -1
  8. data/ext/faiss/utils.h +1 -1
  9. data/lib/faiss/version.rb +1 -1
  10. data/vendor/faiss/faiss/AutoTune.cpp +292 -291
  11. data/vendor/faiss/faiss/AutoTune.h +55 -56
  12. data/vendor/faiss/faiss/Clustering.cpp +365 -194
  13. data/vendor/faiss/faiss/Clustering.h +102 -35
  14. data/vendor/faiss/faiss/IVFlib.cpp +171 -195
  15. data/vendor/faiss/faiss/IVFlib.h +48 -51
  16. data/vendor/faiss/faiss/Index.cpp +85 -103
  17. data/vendor/faiss/faiss/Index.h +54 -48
  18. data/vendor/faiss/faiss/Index2Layer.cpp +126 -224
  19. data/vendor/faiss/faiss/Index2Layer.h +22 -36
  20. data/vendor/faiss/faiss/IndexAdditiveQuantizer.cpp +407 -0
  21. data/vendor/faiss/faiss/IndexAdditiveQuantizer.h +195 -0
  22. data/vendor/faiss/faiss/IndexBinary.cpp +45 -37
  23. data/vendor/faiss/faiss/IndexBinary.h +140 -132
  24. data/vendor/faiss/faiss/IndexBinaryFlat.cpp +73 -53
  25. data/vendor/faiss/faiss/IndexBinaryFlat.h +29 -24
  26. data/vendor/faiss/faiss/IndexBinaryFromFloat.cpp +46 -43
  27. data/vendor/faiss/faiss/IndexBinaryFromFloat.h +16 -15
  28. data/vendor/faiss/faiss/IndexBinaryHNSW.cpp +215 -232
  29. data/vendor/faiss/faiss/IndexBinaryHNSW.h +25 -24
  30. data/vendor/faiss/faiss/IndexBinaryHash.cpp +182 -177
  31. data/vendor/faiss/faiss/IndexBinaryHash.h +41 -34
  32. data/vendor/faiss/faiss/IndexBinaryIVF.cpp +489 -461
  33. data/vendor/faiss/faiss/IndexBinaryIVF.h +97 -68
  34. data/vendor/faiss/faiss/IndexFlat.cpp +115 -176
  35. data/vendor/faiss/faiss/IndexFlat.h +42 -59
  36. data/vendor/faiss/faiss/IndexFlatCodes.cpp +67 -0
  37. data/vendor/faiss/faiss/IndexFlatCodes.h +47 -0
  38. data/vendor/faiss/faiss/IndexHNSW.cpp +372 -348
  39. data/vendor/faiss/faiss/IndexHNSW.h +57 -41
  40. data/vendor/faiss/faiss/IndexIVF.cpp +545 -453
  41. data/vendor/faiss/faiss/IndexIVF.h +169 -118
  42. data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.cpp +316 -0
  43. data/vendor/faiss/faiss/IndexIVFAdditiveQuantizer.h +121 -0
  44. data/vendor/faiss/faiss/IndexIVFFlat.cpp +247 -252
  45. data/vendor/faiss/faiss/IndexIVFFlat.h +48 -51
  46. data/vendor/faiss/faiss/IndexIVFPQ.cpp +459 -517
  47. data/vendor/faiss/faiss/IndexIVFPQ.h +75 -67
  48. data/vendor/faiss/faiss/IndexIVFPQFastScan.cpp +406 -372
  49. data/vendor/faiss/faiss/IndexIVFPQFastScan.h +82 -57
  50. data/vendor/faiss/faiss/IndexIVFPQR.cpp +104 -102
  51. data/vendor/faiss/faiss/IndexIVFPQR.h +33 -28
  52. data/vendor/faiss/faiss/IndexIVFSpectralHash.cpp +163 -150
  53. data/vendor/faiss/faiss/IndexIVFSpectralHash.h +38 -25
  54. data/vendor/faiss/faiss/IndexLSH.cpp +66 -113
  55. data/vendor/faiss/faiss/IndexLSH.h +20 -38
  56. data/vendor/faiss/faiss/IndexLattice.cpp +42 -56
  57. data/vendor/faiss/faiss/IndexLattice.h +11 -16
  58. data/vendor/faiss/faiss/IndexNNDescent.cpp +229 -0
  59. data/vendor/faiss/faiss/IndexNNDescent.h +72 -0
  60. data/vendor/faiss/faiss/IndexNSG.cpp +301 -0
  61. data/vendor/faiss/faiss/IndexNSG.h +85 -0
  62. data/vendor/faiss/faiss/IndexPQ.cpp +387 -495
  63. data/vendor/faiss/faiss/IndexPQ.h +64 -82
  64. data/vendor/faiss/faiss/IndexPQFastScan.cpp +143 -170
  65. data/vendor/faiss/faiss/IndexPQFastScan.h +46 -32
  66. data/vendor/faiss/faiss/IndexPreTransform.cpp +120 -150
  67. data/vendor/faiss/faiss/IndexPreTransform.h +33 -36
  68. data/vendor/faiss/faiss/IndexRefine.cpp +139 -127
  69. data/vendor/faiss/faiss/IndexRefine.h +32 -23
  70. data/vendor/faiss/faiss/IndexReplicas.cpp +147 -153
  71. data/vendor/faiss/faiss/IndexReplicas.h +62 -56
  72. data/vendor/faiss/faiss/IndexScalarQuantizer.cpp +111 -172
  73. data/vendor/faiss/faiss/IndexScalarQuantizer.h +41 -59
  74. data/vendor/faiss/faiss/IndexShards.cpp +256 -240
  75. data/vendor/faiss/faiss/IndexShards.h +85 -73
  76. data/vendor/faiss/faiss/MatrixStats.cpp +112 -97
  77. data/vendor/faiss/faiss/MatrixStats.h +7 -10
  78. data/vendor/faiss/faiss/MetaIndexes.cpp +135 -157
  79. data/vendor/faiss/faiss/MetaIndexes.h +40 -34
  80. data/vendor/faiss/faiss/MetricType.h +7 -7
  81. data/vendor/faiss/faiss/VectorTransform.cpp +654 -475
  82. data/vendor/faiss/faiss/VectorTransform.h +64 -89
  83. data/vendor/faiss/faiss/clone_index.cpp +78 -73
  84. data/vendor/faiss/faiss/clone_index.h +4 -9
  85. data/vendor/faiss/faiss/gpu/GpuAutoTune.cpp +33 -38
  86. data/vendor/faiss/faiss/gpu/GpuAutoTune.h +11 -9
  87. data/vendor/faiss/faiss/gpu/GpuCloner.cpp +198 -171
  88. data/vendor/faiss/faiss/gpu/GpuCloner.h +53 -35
  89. data/vendor/faiss/faiss/gpu/GpuClonerOptions.cpp +12 -14
  90. data/vendor/faiss/faiss/gpu/GpuClonerOptions.h +27 -25
  91. data/vendor/faiss/faiss/gpu/GpuDistance.h +116 -112
  92. data/vendor/faiss/faiss/gpu/GpuFaissAssert.h +1 -2
  93. data/vendor/faiss/faiss/gpu/GpuIcmEncoder.h +60 -0
  94. data/vendor/faiss/faiss/gpu/GpuIndex.h +134 -137
  95. data/vendor/faiss/faiss/gpu/GpuIndexBinaryFlat.h +76 -73
  96. data/vendor/faiss/faiss/gpu/GpuIndexFlat.h +173 -162
  97. data/vendor/faiss/faiss/gpu/GpuIndexIVF.h +67 -64
  98. data/vendor/faiss/faiss/gpu/GpuIndexIVFFlat.h +89 -86
  99. data/vendor/faiss/faiss/gpu/GpuIndexIVFPQ.h +150 -141
  100. data/vendor/faiss/faiss/gpu/GpuIndexIVFScalarQuantizer.h +101 -103
  101. data/vendor/faiss/faiss/gpu/GpuIndicesOptions.h +17 -16
  102. data/vendor/faiss/faiss/gpu/GpuResources.cpp +116 -128
  103. data/vendor/faiss/faiss/gpu/GpuResources.h +182 -186
  104. data/vendor/faiss/faiss/gpu/StandardGpuResources.cpp +433 -422
  105. data/vendor/faiss/faiss/gpu/StandardGpuResources.h +131 -130
  106. data/vendor/faiss/faiss/gpu/impl/InterleavedCodes.cpp +468 -456
  107. data/vendor/faiss/faiss/gpu/impl/InterleavedCodes.h +25 -19
  108. data/vendor/faiss/faiss/gpu/impl/RemapIndices.cpp +22 -20
  109. data/vendor/faiss/faiss/gpu/impl/RemapIndices.h +9 -8
  110. data/vendor/faiss/faiss/gpu/perf/IndexWrapper-inl.h +39 -44
  111. data/vendor/faiss/faiss/gpu/perf/IndexWrapper.h +16 -14
  112. data/vendor/faiss/faiss/gpu/perf/PerfClustering.cpp +77 -71
  113. data/vendor/faiss/faiss/gpu/perf/PerfIVFPQAdd.cpp +109 -88
  114. data/vendor/faiss/faiss/gpu/perf/WriteIndex.cpp +75 -64
  115. data/vendor/faiss/faiss/gpu/test/TestCodePacking.cpp +230 -215
  116. data/vendor/faiss/faiss/gpu/test/TestGpuIndexBinaryFlat.cpp +80 -86
  117. data/vendor/faiss/faiss/gpu/test/TestGpuIndexFlat.cpp +284 -277
  118. data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFFlat.cpp +416 -416
  119. data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFPQ.cpp +611 -517
  120. data/vendor/faiss/faiss/gpu/test/TestGpuIndexIVFScalarQuantizer.cpp +166 -164
  121. data/vendor/faiss/faiss/gpu/test/TestGpuMemoryException.cpp +61 -53
  122. data/vendor/faiss/faiss/gpu/test/TestUtils.cpp +274 -238
  123. data/vendor/faiss/faiss/gpu/test/TestUtils.h +73 -57
  124. data/vendor/faiss/faiss/gpu/test/demo_ivfpq_indexing_gpu.cpp +47 -50
  125. data/vendor/faiss/faiss/gpu/utils/DeviceUtils.h +79 -72
  126. data/vendor/faiss/faiss/gpu/utils/StackDeviceMemory.cpp +140 -146
  127. data/vendor/faiss/faiss/gpu/utils/StackDeviceMemory.h +69 -71
  128. data/vendor/faiss/faiss/gpu/utils/StaticUtils.h +21 -16
  129. data/vendor/faiss/faiss/gpu/utils/Timer.cpp +25 -29
  130. data/vendor/faiss/faiss/gpu/utils/Timer.h +30 -29
  131. data/vendor/faiss/faiss/impl/AdditiveQuantizer.cpp +503 -0
  132. data/vendor/faiss/faiss/impl/AdditiveQuantizer.h +175 -0
  133. data/vendor/faiss/faiss/impl/AuxIndexStructures.cpp +90 -120
  134. data/vendor/faiss/faiss/impl/AuxIndexStructures.h +81 -65
  135. data/vendor/faiss/faiss/impl/FaissAssert.h +73 -58
  136. data/vendor/faiss/faiss/impl/FaissException.cpp +56 -48
  137. data/vendor/faiss/faiss/impl/FaissException.h +41 -29
  138. data/vendor/faiss/faiss/impl/HNSW.cpp +606 -617
  139. data/vendor/faiss/faiss/impl/HNSW.h +179 -200
  140. data/vendor/faiss/faiss/impl/LocalSearchQuantizer.cpp +855 -0
  141. data/vendor/faiss/faiss/impl/LocalSearchQuantizer.h +244 -0
  142. data/vendor/faiss/faiss/impl/NNDescent.cpp +487 -0
  143. data/vendor/faiss/faiss/impl/NNDescent.h +154 -0
  144. data/vendor/faiss/faiss/impl/NSG.cpp +679 -0
  145. data/vendor/faiss/faiss/impl/NSG.h +199 -0
  146. data/vendor/faiss/faiss/impl/PolysemousTraining.cpp +484 -454
  147. data/vendor/faiss/faiss/impl/PolysemousTraining.h +52 -55
  148. data/vendor/faiss/faiss/impl/ProductQuantizer-inl.h +26 -47
  149. data/vendor/faiss/faiss/impl/ProductQuantizer.cpp +469 -459
  150. data/vendor/faiss/faiss/impl/ProductQuantizer.h +76 -87
  151. data/vendor/faiss/faiss/impl/ResidualQuantizer.cpp +758 -0
  152. data/vendor/faiss/faiss/impl/ResidualQuantizer.h +188 -0
  153. data/vendor/faiss/faiss/impl/ResultHandler.h +96 -132
  154. data/vendor/faiss/faiss/impl/ScalarQuantizer.cpp +647 -707
  155. data/vendor/faiss/faiss/impl/ScalarQuantizer.h +48 -46
  156. data/vendor/faiss/faiss/impl/ThreadedIndex-inl.h +129 -131
  157. data/vendor/faiss/faiss/impl/ThreadedIndex.h +61 -55
  158. data/vendor/faiss/faiss/impl/index_read.cpp +631 -480
  159. data/vendor/faiss/faiss/impl/index_write.cpp +547 -407
  160. data/vendor/faiss/faiss/impl/io.cpp +76 -95
  161. data/vendor/faiss/faiss/impl/io.h +31 -41
  162. data/vendor/faiss/faiss/impl/io_macros.h +60 -29
  163. data/vendor/faiss/faiss/impl/kmeans1d.cpp +301 -0
  164. data/vendor/faiss/faiss/impl/kmeans1d.h +48 -0
  165. data/vendor/faiss/faiss/impl/lattice_Zn.cpp +137 -186
  166. data/vendor/faiss/faiss/impl/lattice_Zn.h +40 -51
  167. data/vendor/faiss/faiss/impl/platform_macros.h +29 -8
  168. data/vendor/faiss/faiss/impl/pq4_fast_scan.cpp +77 -124
  169. data/vendor/faiss/faiss/impl/pq4_fast_scan.h +39 -48
  170. data/vendor/faiss/faiss/impl/pq4_fast_scan_search_1.cpp +41 -52
  171. data/vendor/faiss/faiss/impl/pq4_fast_scan_search_qbs.cpp +80 -117
  172. data/vendor/faiss/faiss/impl/simd_result_handlers.h +109 -137
  173. data/vendor/faiss/faiss/index_factory.cpp +619 -397
  174. data/vendor/faiss/faiss/index_factory.h +8 -6
  175. data/vendor/faiss/faiss/index_io.h +23 -26
  176. data/vendor/faiss/faiss/invlists/BlockInvertedLists.cpp +67 -75
  177. data/vendor/faiss/faiss/invlists/BlockInvertedLists.h +22 -24
  178. data/vendor/faiss/faiss/invlists/DirectMap.cpp +96 -112
  179. data/vendor/faiss/faiss/invlists/DirectMap.h +29 -33
  180. data/vendor/faiss/faiss/invlists/InvertedLists.cpp +307 -364
  181. data/vendor/faiss/faiss/invlists/InvertedLists.h +151 -151
  182. data/vendor/faiss/faiss/invlists/InvertedListsIOHook.cpp +29 -34
  183. data/vendor/faiss/faiss/invlists/InvertedListsIOHook.h +17 -18
  184. data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.cpp +257 -293
  185. data/vendor/faiss/faiss/invlists/OnDiskInvertedLists.h +50 -45
  186. data/vendor/faiss/faiss/python/python_callbacks.cpp +23 -26
  187. data/vendor/faiss/faiss/python/python_callbacks.h +9 -16
  188. data/vendor/faiss/faiss/utils/AlignedTable.h +79 -44
  189. data/vendor/faiss/faiss/utils/Heap.cpp +40 -48
  190. data/vendor/faiss/faiss/utils/Heap.h +186 -209
  191. data/vendor/faiss/faiss/utils/WorkerThread.cpp +67 -76
  192. data/vendor/faiss/faiss/utils/WorkerThread.h +32 -33
  193. data/vendor/faiss/faiss/utils/distances.cpp +305 -312
  194. data/vendor/faiss/faiss/utils/distances.h +170 -122
  195. data/vendor/faiss/faiss/utils/distances_simd.cpp +498 -508
  196. data/vendor/faiss/faiss/utils/extra_distances-inl.h +117 -0
  197. data/vendor/faiss/faiss/utils/extra_distances.cpp +113 -232
  198. data/vendor/faiss/faiss/utils/extra_distances.h +30 -29
  199. data/vendor/faiss/faiss/utils/hamming-inl.h +260 -209
  200. data/vendor/faiss/faiss/utils/hamming.cpp +375 -469
  201. data/vendor/faiss/faiss/utils/hamming.h +62 -85
  202. data/vendor/faiss/faiss/utils/ordered_key_value.h +16 -18
  203. data/vendor/faiss/faiss/utils/partitioning.cpp +393 -318
  204. data/vendor/faiss/faiss/utils/partitioning.h +26 -21
  205. data/vendor/faiss/faiss/utils/quantize_lut.cpp +78 -66
  206. data/vendor/faiss/faiss/utils/quantize_lut.h +22 -20
  207. data/vendor/faiss/faiss/utils/random.cpp +39 -63
  208. data/vendor/faiss/faiss/utils/random.h +13 -16
  209. data/vendor/faiss/faiss/utils/simdlib.h +4 -2
  210. data/vendor/faiss/faiss/utils/simdlib_avx2.h +88 -85
  211. data/vendor/faiss/faiss/utils/simdlib_emulated.h +226 -165
  212. data/vendor/faiss/faiss/utils/simdlib_neon.h +832 -0
  213. data/vendor/faiss/faiss/utils/utils.cpp +304 -287
  214. data/vendor/faiss/faiss/utils/utils.h +54 -49
  215. metadata +29 -4
@@ -5,12 +5,11 @@
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
  #include <sstream>
16
15
  #include <vector>
@@ -20,93 +19,109 @@ constexpr float kF16MaxRelErr = 0.07f;
20
19
  constexpr float kF32MaxRelErr = 6e-3f;
21
20
 
22
21
  struct TestFlatOptions {
23
- TestFlatOptions()
24
- : metric(faiss::MetricType::METRIC_L2),
25
- metricArg(0),
26
- useFloat16(false),
27
- useTransposed(false),
28
- numVecsOverride(-1),
29
- numQueriesOverride(-1),
30
- kOverride(-1),
31
- dimOverride(-1) {
32
- }
33
-
34
- faiss::MetricType metric;
35
- float metricArg;
36
-
37
- bool useFloat16;
38
- bool useTransposed;
39
- int numVecsOverride;
40
- int numQueriesOverride;
41
- int kOverride;
42
- int dimOverride;
22
+ TestFlatOptions()
23
+ : metric(faiss::MetricType::METRIC_L2),
24
+ metricArg(0),
25
+ useFloat16(false),
26
+ useTransposed(false),
27
+ numVecsOverride(-1),
28
+ numQueriesOverride(-1),
29
+ kOverride(-1),
30
+ dimOverride(-1) {}
31
+
32
+ faiss::MetricType metric;
33
+ float metricArg;
34
+
35
+ bool useFloat16;
36
+ bool useTransposed;
37
+ int numVecsOverride;
38
+ int numQueriesOverride;
39
+ int kOverride;
40
+ int dimOverride;
43
41
  };
44
42
 
45
43
  void testFlat(const TestFlatOptions& opt) {
46
- int numVecs = opt.numVecsOverride > 0 ?
47
- opt.numVecsOverride : faiss::gpu::randVal(1000, 5000);
48
- int dim = opt.dimOverride > 0 ?
49
- opt.dimOverride : faiss::gpu::randVal(50, 800);
50
- int numQuery = opt.numQueriesOverride > 0 ?
51
- opt.numQueriesOverride : faiss::gpu::randVal(1, 512);
52
-
53
- // Due to loss of precision in a float16 accumulator, for large k,
54
- // the number of differences is pretty huge. Restrict ourselves to a
55
- // fairly small `k` for float16
56
- int k = opt.useFloat16 ?
57
- std::min(faiss::gpu::randVal(1, 50), numVecs) :
58
- std::min(faiss::gpu::randVal(1, faiss::gpu::getMaxKSelection()), numVecs);
59
- if (opt.kOverride > 0) {
60
- k = opt.kOverride;
61
- }
62
-
63
- faiss::IndexFlat cpuIndex(dim, opt.metric);
64
- cpuIndex.metric_arg = opt.metricArg;
65
-
66
- // Construct on a random device to test multi-device, if we have
67
- // multiple devices
68
- int device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
69
-
70
- faiss::gpu::StandardGpuResources res;
71
- res.noTempMemory();
72
-
73
- faiss::gpu::GpuIndexFlatConfig config;
74
- config.device = device;
75
- config.useFloat16 = opt.useFloat16;
76
- config.storeTransposed = opt.useTransposed;
77
-
78
- faiss::gpu::GpuIndexFlat gpuIndex(&res, dim, opt.metric, config);
79
- gpuIndex.metric_arg = opt.metricArg;
80
-
81
- std::vector<float> vecs = faiss::gpu::randVecs(numVecs, dim);
82
- cpuIndex.add(numVecs, vecs.data());
83
- gpuIndex.add(numVecs, vecs.data());
84
-
85
- std::stringstream str;
86
- str << "metric " << opt.metric
87
- << " marg " << opt.metricArg
88
- << " numVecs " << numVecs
89
- << " dim " << dim
90
- << " useFloat16 " << opt.useFloat16
91
- << " transposed " << opt.useTransposed
92
- << " numQuery " << numQuery
93
- << " k " << k;
94
-
95
- // To some extent, we depend upon the relative error for the test
96
- // for float16
97
- faiss::gpu::compareIndices(cpuIndex, gpuIndex, numQuery, dim, k, str.str(),
98
- opt.useFloat16 ? kF16MaxRelErr : kF32MaxRelErr,
99
- // FIXME: the fp16 bounds are
100
- // useless when math (the accumulator) is
101
- // in fp16. Figure out another way to test
102
- opt.useFloat16 ? 0.99f : 0.1f,
103
- opt.useFloat16 ? 0.65f : 0.015f);
44
+ int numVecs = opt.numVecsOverride > 0 ? opt.numVecsOverride
45
+ : faiss::gpu::randVal(1000, 5000);
46
+ int dim = opt.dimOverride > 0 ? opt.dimOverride
47
+ : faiss::gpu::randVal(50, 800);
48
+ int numQuery = opt.numQueriesOverride > 0 ? opt.numQueriesOverride
49
+ : faiss::gpu::randVal(1, 512);
50
+
51
+ // Due to loss of precision in a float16 accumulator, for large k,
52
+ // the number of differences is pretty huge. Restrict ourselves to a
53
+ // fairly small `k` for float16
54
+ int k = opt.useFloat16
55
+ ? std::min(faiss::gpu::randVal(1, 50), numVecs)
56
+ : std::min(
57
+ faiss::gpu::randVal(1, faiss::gpu::getMaxKSelection()),
58
+ numVecs);
59
+ if (opt.kOverride > 0) {
60
+ k = opt.kOverride;
61
+ }
62
+
63
+ faiss::IndexFlat cpuIndex(dim, opt.metric);
64
+ cpuIndex.metric_arg = opt.metricArg;
65
+
66
+ // Construct on a random device to test multi-device, if we have
67
+ // multiple devices
68
+ int device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
69
+
70
+ faiss::gpu::StandardGpuResources res;
71
+ res.noTempMemory();
72
+
73
+ faiss::gpu::GpuIndexFlatConfig config;
74
+ config.device = device;
75
+ config.useFloat16 = opt.useFloat16;
76
+ config.storeTransposed = opt.useTransposed;
77
+
78
+ faiss::gpu::GpuIndexFlat gpuIndex(&res, dim, opt.metric, config);
79
+ gpuIndex.metric_arg = opt.metricArg;
80
+
81
+ std::vector<float> vecs = faiss::gpu::randVecs(numVecs, dim);
82
+ cpuIndex.add(numVecs, vecs.data());
83
+ gpuIndex.add(numVecs, vecs.data());
84
+
85
+ std::stringstream str;
86
+ str << "metric " << opt.metric << " marg " << opt.metricArg << " numVecs "
87
+ << numVecs << " dim " << dim << " useFloat16 " << opt.useFloat16
88
+ << " transposed " << opt.useTransposed << " numQuery " << numQuery
89
+ << " k " << k;
90
+
91
+ // To some extent, we depend upon the relative error for the test
92
+ // for float16
93
+ faiss::gpu::compareIndices(
94
+ cpuIndex,
95
+ gpuIndex,
96
+ numQuery,
97
+ dim,
98
+ k,
99
+ str.str(),
100
+ opt.useFloat16 ? kF16MaxRelErr : kF32MaxRelErr,
101
+ // FIXME: the fp16 bounds are
102
+ // useless when math (the accumulator) is
103
+ // in fp16. Figure out another way to test
104
+ opt.useFloat16 ? 0.99f : 0.1f,
105
+ opt.useFloat16 ? 0.65f : 0.015f);
104
106
  }
105
107
 
106
108
  TEST(TestGpuIndexFlat, IP_Float32) {
107
- for (int tries = 0; tries < 3; ++tries) {
109
+ for (int tries = 0; tries < 3; ++tries) {
110
+ TestFlatOptions opt;
111
+ opt.metric = faiss::MetricType::METRIC_INNER_PRODUCT;
112
+ opt.useFloat16 = false;
113
+ opt.useTransposed = false;
114
+
115
+ testFlat(opt);
116
+
117
+ opt.useTransposed = true;
118
+ testFlat(opt);
119
+ }
120
+ }
121
+
122
+ TEST(TestGpuIndexFlat, L1_Float32) {
108
123
  TestFlatOptions opt;
109
- opt.metric = faiss::MetricType::METRIC_INNER_PRODUCT;
124
+ opt.metric = faiss::MetricType::METRIC_L1;
110
125
  opt.useFloat16 = false;
111
126
  opt.useTransposed = false;
112
127
 
@@ -114,280 +129,272 @@ TEST(TestGpuIndexFlat, IP_Float32) {
114
129
 
115
130
  opt.useTransposed = true;
116
131
  testFlat(opt);
117
- }
118
- }
119
-
120
- TEST(TestGpuIndexFlat, L1_Float32) {
121
- TestFlatOptions opt;
122
- opt.metric = faiss::MetricType::METRIC_L1;
123
- opt.useFloat16 = false;
124
- opt.useTransposed = false;
125
-
126
- testFlat(opt);
127
-
128
- opt.useTransposed = true;
129
- testFlat(opt);
130
132
  }
131
133
 
132
134
  TEST(TestGpuIndexFlat, Lp_Float32) {
133
- TestFlatOptions opt;
134
- opt.metric = faiss::MetricType::METRIC_Lp;
135
- opt.metricArg = 5;
136
- opt.useFloat16 = false;
137
- opt.useTransposed = false;
135
+ TestFlatOptions opt;
136
+ opt.metric = faiss::MetricType::METRIC_Lp;
137
+ opt.metricArg = 5;
138
+ opt.useFloat16 = false;
139
+ opt.useTransposed = false;
138
140
 
139
- testFlat(opt);
141
+ testFlat(opt);
140
142
 
141
- // Don't bother testing the transposed version, the L1 test should be good
142
- // enough for that
143
+ // Don't bother testing the transposed version, the L1 test should be good
144
+ // enough for that
143
145
  }
144
146
 
145
147
  TEST(TestGpuIndexFlat, L2_Float32) {
146
- for (int tries = 0; tries < 3; ++tries) {
147
- TestFlatOptions opt;
148
- opt.metric = faiss::MetricType::METRIC_L2;
148
+ for (int tries = 0; tries < 3; ++tries) {
149
+ TestFlatOptions opt;
150
+ opt.metric = faiss::MetricType::METRIC_L2;
149
151
 
150
- opt.useFloat16 = false;
151
- opt.useTransposed = false;
152
+ opt.useFloat16 = false;
153
+ opt.useTransposed = false;
152
154
 
153
- testFlat(opt);
155
+ testFlat(opt);
154
156
 
155
- opt.useTransposed = true;
156
- testFlat(opt);
157
- }
157
+ opt.useTransposed = true;
158
+ testFlat(opt);
159
+ }
158
160
  }
159
161
 
160
162
  // test specialized k == 1 codepath
161
163
  TEST(TestGpuIndexFlat, L2_Float32_K1) {
162
- for (int tries = 0; tries < 3; ++tries) {
163
- TestFlatOptions opt;
164
- opt.metric = faiss::MetricType::METRIC_L2;
165
- opt.useFloat16 = false;
166
- opt.useTransposed = false;
167
- opt.kOverride = 1;
168
-
169
- testFlat(opt);
170
- }
164
+ for (int tries = 0; tries < 3; ++tries) {
165
+ TestFlatOptions opt;
166
+ opt.metric = faiss::MetricType::METRIC_L2;
167
+ opt.useFloat16 = false;
168
+ opt.useTransposed = false;
169
+ opt.kOverride = 1;
170
+
171
+ testFlat(opt);
172
+ }
171
173
  }
172
174
 
173
175
  TEST(TestGpuIndexFlat, IP_Float16) {
174
- for (int tries = 0; tries < 3; ++tries) {
175
- TestFlatOptions opt;
176
- opt.metric = faiss::MetricType::METRIC_INNER_PRODUCT;
177
- opt.useFloat16 = true;
178
- opt.useTransposed = false;
176
+ for (int tries = 0; tries < 3; ++tries) {
177
+ TestFlatOptions opt;
178
+ opt.metric = faiss::MetricType::METRIC_INNER_PRODUCT;
179
+ opt.useFloat16 = true;
180
+ opt.useTransposed = false;
179
181
 
180
- testFlat(opt);
182
+ testFlat(opt);
181
183
 
182
- opt.useTransposed = true;
183
- testFlat(opt);
184
- }
184
+ opt.useTransposed = true;
185
+ testFlat(opt);
186
+ }
185
187
  }
186
188
 
187
189
  TEST(TestGpuIndexFlat, L2_Float16) {
188
- for (int tries = 0; tries < 3; ++tries) {
189
- TestFlatOptions opt;
190
- opt.metric = faiss::MetricType::METRIC_L2;
191
- opt.useFloat16 = true;
192
- opt.useTransposed = false;
190
+ for (int tries = 0; tries < 3; ++tries) {
191
+ TestFlatOptions opt;
192
+ opt.metric = faiss::MetricType::METRIC_L2;
193
+ opt.useFloat16 = true;
194
+ opt.useTransposed = false;
193
195
 
194
- testFlat(opt);
196
+ testFlat(opt);
195
197
 
196
- opt.useTransposed = true;
197
- testFlat(opt);
198
- }
198
+ opt.useTransposed = true;
199
+ testFlat(opt);
200
+ }
199
201
  }
200
202
 
201
203
  // test specialized k == 1 codepath
202
204
  TEST(TestGpuIndexFlat, L2_Float16_K1) {
203
- for (int tries = 0; tries < 3; ++tries) {
204
- TestFlatOptions opt;
205
- opt.metric = faiss::MetricType::METRIC_L2;
206
- opt.useFloat16 = true;
207
- opt.useTransposed = false;
208
- opt.kOverride = 1;
209
-
210
- testFlat(opt);
211
- }
205
+ for (int tries = 0; tries < 3; ++tries) {
206
+ TestFlatOptions opt;
207
+ opt.metric = faiss::MetricType::METRIC_L2;
208
+ opt.useFloat16 = true;
209
+ opt.useTransposed = false;
210
+ opt.kOverride = 1;
211
+
212
+ testFlat(opt);
213
+ }
212
214
  }
213
215
 
214
216
  // test tiling along a huge vector set
215
217
  TEST(TestGpuIndexFlat, L2_Tiling) {
216
- for (int tries = 0; tries < 2; ++tries) {
217
- TestFlatOptions opt;
218
- opt.metric = faiss::MetricType::METRIC_L2;
219
- opt.useFloat16 = false;
220
- opt.useTransposed = false;
221
- opt.numVecsOverride = 1000000;
222
-
223
- // keep the rest of the problem reasonably small
224
- opt.numQueriesOverride = 4;
225
- opt.dimOverride = 64;
226
- opt.kOverride = 64;
227
-
228
- testFlat(opt);
229
- }
218
+ for (int tries = 0; tries < 2; ++tries) {
219
+ TestFlatOptions opt;
220
+ opt.metric = faiss::MetricType::METRIC_L2;
221
+ opt.useFloat16 = false;
222
+ opt.useTransposed = false;
223
+ opt.numVecsOverride = 1000000;
224
+
225
+ // keep the rest of the problem reasonably small
226
+ opt.numQueriesOverride = 4;
227
+ opt.dimOverride = 64;
228
+ opt.kOverride = 64;
229
+
230
+ testFlat(opt);
231
+ }
230
232
  }
231
233
 
232
234
  TEST(TestGpuIndexFlat, QueryEmpty) {
233
- faiss::gpu::StandardGpuResources res;
234
- res.noTempMemory();
235
+ faiss::gpu::StandardGpuResources res;
236
+ res.noTempMemory();
235
237
 
236
- faiss::gpu::GpuIndexFlatConfig config;
237
- config.device = 0;
238
- config.useFloat16 = false;
239
- config.storeTransposed = false;
238
+ faiss::gpu::GpuIndexFlatConfig config;
239
+ config.device = 0;
240
+ config.useFloat16 = false;
241
+ config.storeTransposed = false;
240
242
 
241
- int dim = 128;
242
- faiss::gpu::GpuIndexFlatL2 gpuIndex(&res, dim, config);
243
+ int dim = 128;
244
+ faiss::gpu::GpuIndexFlatL2 gpuIndex(&res, dim, config);
243
245
 
244
- // Querying an empty index should not blow up, and just return
245
- // (FLT_MAX, -1)
246
- int numQuery = 10;
247
- int k = 50;
248
- std::vector<float> queries(numQuery * dim, 1.0f);
246
+ // Querying an empty index should not blow up, and just return
247
+ // (FLT_MAX, -1)
248
+ int numQuery = 10;
249
+ int k = 50;
250
+ std::vector<float> queries(numQuery * dim, 1.0f);
249
251
 
250
- std::vector<float> dist(numQuery * k, 0);
251
- std::vector<faiss::Index::idx_t> ind(numQuery * k);
252
+ std::vector<float> dist(numQuery * k, 0);
253
+ std::vector<faiss::Index::idx_t> ind(numQuery * k);
252
254
 
253
- gpuIndex.search(numQuery, queries.data(), k, dist.data(), ind.data());
255
+ gpuIndex.search(numQuery, queries.data(), k, dist.data(), ind.data());
254
256
 
255
- for (auto d : dist) {
256
- EXPECT_EQ(d, std::numeric_limits<float>::max());
257
- }
257
+ for (auto d : dist) {
258
+ EXPECT_EQ(d, std::numeric_limits<float>::max());
259
+ }
258
260
 
259
- for (auto i : ind) {
260
- EXPECT_EQ(i, -1);
261
- }
261
+ for (auto i : ind) {
262
+ EXPECT_EQ(i, -1);
263
+ }
262
264
  }
263
265
 
264
266
  TEST(TestGpuIndexFlat, CopyFrom) {
265
- int numVecs = faiss::gpu::randVal(100, 200);
266
- int dim = faiss::gpu::randVal(1, 1000);
267
+ int numVecs = faiss::gpu::randVal(100, 200);
268
+ int dim = faiss::gpu::randVal(1, 1000);
267
269
 
268
- faiss::IndexFlatL2 cpuIndex(dim);
270
+ faiss::IndexFlatL2 cpuIndex(dim);
269
271
 
270
- std::vector<float> vecs = faiss::gpu::randVecs(numVecs, dim);
271
- cpuIndex.add(numVecs, vecs.data());
272
+ std::vector<float> vecs = faiss::gpu::randVecs(numVecs, dim);
273
+ cpuIndex.add(numVecs, vecs.data());
272
274
 
273
- faiss::gpu::StandardGpuResources res;
274
- res.noTempMemory();
275
+ faiss::gpu::StandardGpuResources res;
276
+ res.noTempMemory();
275
277
 
276
- // Fill with garbage values
277
- int device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
278
+ // Fill with garbage values
279
+ int device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
278
280
 
279
- faiss::gpu::GpuIndexFlatConfig config;
280
- config.device = device;
281
- config.useFloat16 = false;
282
- config.storeTransposed = false;
281
+ faiss::gpu::GpuIndexFlatConfig config;
282
+ config.device = device;
283
+ config.useFloat16 = false;
284
+ config.storeTransposed = false;
283
285
 
284
- faiss::gpu::GpuIndexFlatL2 gpuIndex(&res, 2000, config);
285
- gpuIndex.copyFrom(&cpuIndex);
286
+ faiss::gpu::GpuIndexFlatL2 gpuIndex(&res, 2000, config);
287
+ gpuIndex.copyFrom(&cpuIndex);
286
288
 
287
- EXPECT_EQ(cpuIndex.ntotal, gpuIndex.ntotal);
288
- EXPECT_EQ(gpuIndex.ntotal, numVecs);
289
+ EXPECT_EQ(cpuIndex.ntotal, gpuIndex.ntotal);
290
+ EXPECT_EQ(gpuIndex.ntotal, numVecs);
289
291
 
290
- EXPECT_EQ(cpuIndex.d, gpuIndex.d);
291
- EXPECT_EQ(cpuIndex.d, dim);
292
+ EXPECT_EQ(cpuIndex.d, gpuIndex.d);
293
+ EXPECT_EQ(cpuIndex.d, dim);
292
294
 
293
- int idx = faiss::gpu::randVal(0, numVecs - 1);
295
+ int idx = faiss::gpu::randVal(0, numVecs - 1);
294
296
 
295
- std::vector<float> gpuVals(dim);
296
- gpuIndex.reconstruct(idx, gpuVals.data());
297
+ std::vector<float> gpuVals(dim);
298
+ gpuIndex.reconstruct(idx, gpuVals.data());
297
299
 
298
- std::vector<float> cpuVals(dim);
299
- cpuIndex.reconstruct(idx, cpuVals.data());
300
+ std::vector<float> cpuVals(dim);
301
+ cpuIndex.reconstruct(idx, cpuVals.data());
300
302
 
301
- EXPECT_EQ(gpuVals, cpuVals);
303
+ EXPECT_EQ(gpuVals, cpuVals);
302
304
  }
303
305
 
304
306
  TEST(TestGpuIndexFlat, CopyTo) {
305
- faiss::gpu::StandardGpuResources res;
306
- res.noTempMemory();
307
+ faiss::gpu::StandardGpuResources res;
308
+ res.noTempMemory();
307
309
 
308
- int numVecs = faiss::gpu::randVal(100, 200);
309
- int dim = faiss::gpu::randVal(1, 1000);
310
+ int numVecs = faiss::gpu::randVal(100, 200);
311
+ int dim = faiss::gpu::randVal(1, 1000);
310
312
 
311
- int device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
313
+ int device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
312
314
 
313
- faiss::gpu::GpuIndexFlatConfig config;
314
- config.device = device;
315
- config.useFloat16 = false;
316
- config.storeTransposed = false;
315
+ faiss::gpu::GpuIndexFlatConfig config;
316
+ config.device = device;
317
+ config.useFloat16 = false;
318
+ config.storeTransposed = false;
317
319
 
318
- faiss::gpu::GpuIndexFlatL2 gpuIndex(&res, dim, config);
320
+ faiss::gpu::GpuIndexFlatL2 gpuIndex(&res, dim, config);
319
321
 
320
- std::vector<float> vecs = faiss::gpu::randVecs(numVecs, dim);
321
- gpuIndex.add(numVecs, vecs.data());
322
+ std::vector<float> vecs = faiss::gpu::randVecs(numVecs, dim);
323
+ gpuIndex.add(numVecs, vecs.data());
322
324
 
323
- // Fill with garbage values
324
- faiss::IndexFlatL2 cpuIndex(2000);
325
- gpuIndex.copyTo(&cpuIndex);
325
+ // Fill with garbage values
326
+ faiss::IndexFlatL2 cpuIndex(2000);
327
+ gpuIndex.copyTo(&cpuIndex);
326
328
 
327
- EXPECT_EQ(cpuIndex.ntotal, gpuIndex.ntotal);
328
- EXPECT_EQ(gpuIndex.ntotal, numVecs);
329
+ EXPECT_EQ(cpuIndex.ntotal, gpuIndex.ntotal);
330
+ EXPECT_EQ(gpuIndex.ntotal, numVecs);
329
331
 
330
- EXPECT_EQ(cpuIndex.d, gpuIndex.d);
331
- EXPECT_EQ(cpuIndex.d, dim);
332
+ EXPECT_EQ(cpuIndex.d, gpuIndex.d);
333
+ EXPECT_EQ(cpuIndex.d, dim);
332
334
 
333
- int idx = faiss::gpu::randVal(0, numVecs - 1);
335
+ int idx = faiss::gpu::randVal(0, numVecs - 1);
334
336
 
335
- std::vector<float> gpuVals(dim);
336
- gpuIndex.reconstruct(idx, gpuVals.data());
337
+ std::vector<float> gpuVals(dim);
338
+ gpuIndex.reconstruct(idx, gpuVals.data());
337
339
 
338
- std::vector<float> cpuVals(dim);
339
- cpuIndex.reconstruct(idx, cpuVals.data());
340
+ std::vector<float> cpuVals(dim);
341
+ cpuIndex.reconstruct(idx, cpuVals.data());
340
342
 
341
- EXPECT_EQ(gpuVals, cpuVals);
343
+ EXPECT_EQ(gpuVals, cpuVals);
342
344
  }
343
345
 
344
346
  TEST(TestGpuIndexFlat, UnifiedMemory) {
345
- // Construct on a random device to test multi-device, if we have
346
- // multiple devices
347
- int device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
348
-
349
- if (!faiss::gpu::getFullUnifiedMemSupport(device)) {
350
- return;
351
- }
352
-
353
- int dim = 256;
354
-
355
- // FIXME: GpuIndexFlat doesn't support > 2^31 (vecs * dims) due to
356
- // kernel indexing, so we can't test unified memory for memory
357
- // oversubscription.
358
- size_t numVecs = 50000;
359
- int numQuery = 10;
360
- int k = 10;
361
-
362
- faiss::IndexFlatL2 cpuIndexL2(dim);
363
-
364
- faiss::gpu::StandardGpuResources res;
365
- res.noTempMemory();
366
-
367
- faiss::gpu::GpuIndexFlatConfig config;
368
- config.device = device;
369
- config.memorySpace = faiss::gpu::MemorySpace::Unified;
370
-
371
- faiss::gpu::GpuIndexFlatL2 gpuIndexL2(&res, dim, config);
372
-
373
- std::vector<float> vecs = faiss::gpu::randVecs(numVecs, dim);
374
- cpuIndexL2.add(numVecs, vecs.data());
375
- gpuIndexL2.add(numVecs, vecs.data());
376
-
377
- // To some extent, we depend upon the relative error for the test
378
- // for float16
379
- faiss::gpu::compareIndices(cpuIndexL2, gpuIndexL2,
380
- numQuery, dim, k, "Unified Memory",
381
- kF32MaxRelErr,
382
- 0.1f,
383
- 0.015f);
347
+ // Construct on a random device to test multi-device, if we have
348
+ // multiple devices
349
+ int device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
350
+
351
+ if (!faiss::gpu::getFullUnifiedMemSupport(device)) {
352
+ return;
353
+ }
354
+
355
+ int dim = 256;
356
+
357
+ // FIXME: GpuIndexFlat doesn't support > 2^31 (vecs * dims) due to
358
+ // kernel indexing, so we can't test unified memory for memory
359
+ // oversubscription.
360
+ size_t numVecs = 50000;
361
+ int numQuery = 10;
362
+ int k = 10;
363
+
364
+ faiss::IndexFlatL2 cpuIndexL2(dim);
365
+
366
+ faiss::gpu::StandardGpuResources res;
367
+ res.noTempMemory();
368
+
369
+ faiss::gpu::GpuIndexFlatConfig config;
370
+ config.device = device;
371
+ config.memorySpace = faiss::gpu::MemorySpace::Unified;
372
+
373
+ faiss::gpu::GpuIndexFlatL2 gpuIndexL2(&res, dim, config);
374
+
375
+ std::vector<float> vecs = faiss::gpu::randVecs(numVecs, dim);
376
+ cpuIndexL2.add(numVecs, vecs.data());
377
+ gpuIndexL2.add(numVecs, vecs.data());
378
+
379
+ // To some extent, we depend upon the relative error for the test
380
+ // for float16
381
+ faiss::gpu::compareIndices(
382
+ cpuIndexL2,
383
+ gpuIndexL2,
384
+ numQuery,
385
+ dim,
386
+ k,
387
+ "Unified Memory",
388
+ kF32MaxRelErr,
389
+ 0.1f,
390
+ 0.015f);
384
391
  }
385
392
 
386
393
  int main(int argc, char** argv) {
387
- testing::InitGoogleTest(&argc, argv);
394
+ testing::InitGoogleTest(&argc, argv);
388
395
 
389
- // just run with a fixed test seed
390
- faiss::gpu::setTestSeed(100);
396
+ // just run with a fixed test seed
397
+ faiss::gpu::setTestSeed(100);
391
398
 
392
- return RUN_ALL_TESTS();
399
+ return RUN_ALL_TESTS();
393
400
  }