12 #include "../../IndexFlat.h"
13 #include "../../IndexIVFPQ.h"
14 #include "../GpuIndexIVFPQ.h"
15 #include "../StandardGpuResources.h"
16 #include "../utils/DeviceUtils.h"
17 #include "../test/TestUtils.h"
19 #include <gtest/gtest.h>
24 void pickEncoding(
int& codes,
int& dim) {
25 std::vector<int> codeSizes{3, 4, 8, 12, 16, 20, 24, 28, 32, 40, 48, 56, 64};
26 std::vector<int> dimSizes{4, 8, 16, 32};
28 codes = codeSizes[faiss::gpu::randVal(0, codeSizes.size() - 1)];
31 dim = codes * dimSizes[faiss::gpu::randVal(0, dimSizes.size() - 1)];
35 if (dim < 512 && dim >= 64) {
43 numAdd = faiss::gpu::randVal(10000, 30000);
44 numCentroids = std::sqrt((
float) numAdd);
45 numTrain = numCentroids * 40;
47 pickEncoding(codes, dim);
49 bitsPerCode = faiss::gpu::randVal(3, 8);
50 nprobe = std::min(faiss::gpu::randVal(40, 1000), numCentroids);
51 numQuery = faiss::gpu::randVal(32, 256);
52 k = std::min(faiss::gpu::randVal(10, 50), numAdd / 40);
53 usePrecomputed = faiss::gpu::randBool();
54 indicesOpt = faiss::gpu::randSelect({
55 faiss::gpu::INDICES_CPU,
56 faiss::gpu::INDICES_32_BIT,
57 faiss::gpu::INDICES_64_BIT});
62 useFloat16 = faiss::gpu::randBool();
65 device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
68 std::string toString()
const {
69 std::stringstream str;
70 str <<
"IVFPQ device " << device
71 <<
" numVecs " << numAdd
73 <<
" numCentroids " << numCentroids
75 <<
" bitsPerCode " << bitsPerCode
76 <<
" nprobe " << nprobe
77 <<
" numQuery " << numQuery
79 <<
" usePrecomputed " << usePrecomputed
80 <<
" indicesOpt " << indicesOpt
81 <<
" useFloat16 " << useFloat16;
86 float getCompareEpsilon()
const {
90 float getPctMaxDiff1()
const {
91 return useFloat16 ? 0.30f : 0.10f;
94 float getPctMaxDiffN()
const {
95 return useFloat16 ? 0.05f : 0.015f;
108 faiss::gpu::IndicesOptions indicesOpt;
113 TEST(TestGpuIndexIVFPQ, Query) {
114 for (
int tries = 0; tries < 5; ++tries) {
115 faiss::gpu::newTestSeed();
119 std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
120 std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
124 opt.codes, opt.bitsPerCode);
125 cpuIndex.nprobe = opt.nprobe;
126 cpuIndex.train(opt.numTrain, trainVecs.data());
127 cpuIndex.add(opt.numAdd, addVecs.data());
136 gpuIndex.setPrecomputedCodes(opt.usePrecomputed);
137 gpuIndex.setNumProbes(opt.nprobe);
139 faiss::gpu::compareIndices(cpuIndex, gpuIndex,
140 opt.numQuery, opt.dim, opt.k, opt.toString(),
141 opt.getCompareEpsilon(),
142 opt.getPctMaxDiff1(),
143 opt.getPctMaxDiffN());
147 TEST(TestGpuIndexIVFPQ, Add) {
148 for (
int tries = 0; tries < 5; ++tries) {
149 faiss::gpu::newTestSeed();
153 std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
154 std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
158 opt.codes, opt.bitsPerCode);
159 cpuIndex.nprobe = opt.nprobe;
160 cpuIndex.train(opt.numTrain, trainVecs.data());
169 gpuIndex.setPrecomputedCodes(opt.usePrecomputed);
170 gpuIndex.setNumProbes(opt.nprobe);
172 gpuIndex.add(opt.numAdd, addVecs.data());
173 cpuIndex.add(opt.numAdd, addVecs.data());
175 faiss::gpu::compareIndices(cpuIndex, gpuIndex,
176 opt.numQuery, opt.dim, opt.k, opt.toString(),
177 opt.getCompareEpsilon(),
178 opt.getPctMaxDiff1(),
179 opt.getPctMaxDiffN());
183 TEST(TestGpuIndexIVFPQ, CopyTo) {
184 faiss::gpu::newTestSeed();
187 std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
188 std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
202 gpuIndex.setNumProbes(opt.nprobe);
203 gpuIndex.train(opt.numTrain, trainVecs.data());
204 gpuIndex.add(opt.numAdd, addVecs.data());
210 gpuIndex.copyTo(&cpuIndex);
212 EXPECT_EQ(cpuIndex.ntotal, gpuIndex.ntotal);
213 EXPECT_EQ(gpuIndex.ntotal, opt.numAdd);
215 EXPECT_EQ(cpuIndex.d, gpuIndex.d);
216 EXPECT_EQ(cpuIndex.d, opt.dim);
217 EXPECT_EQ(cpuIndex.nlist, gpuIndex.getNumLists());
218 EXPECT_EQ(cpuIndex.nprobe, gpuIndex.getNumProbes());
219 EXPECT_EQ(cpuIndex.pq.M, gpuIndex.getNumSubQuantizers());
220 EXPECT_EQ(gpuIndex.getNumSubQuantizers(), opt.codes);
221 EXPECT_EQ(cpuIndex.pq.nbits, gpuIndex.getBitsPerCode());
222 EXPECT_EQ(gpuIndex.getBitsPerCode(), opt.bitsPerCode);
225 faiss::gpu::compareIndices(cpuIndex, gpuIndex,
226 opt.numQuery, opt.dim, opt.k, opt.toString(),
227 opt.getCompareEpsilon(),
228 opt.getPctMaxDiff1(),
229 opt.getPctMaxDiffN());
232 TEST(TestGpuIndexIVFPQ, CopyFrom) {
233 faiss::gpu::newTestSeed();
236 std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
237 std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
241 opt.codes, opt.bitsPerCode);
242 cpuIndex.nprobe = opt.nprobe;
243 cpuIndex.train(opt.numTrain, trainVecs.data());
244 cpuIndex.add(opt.numAdd, addVecs.data());
251 1, 1, 1, 1,
false, opt.indicesOpt,
252 opt.useFloat16, faiss::METRIC_L2);
253 gpuIndex.setNumProbes(1);
255 gpuIndex.copyFrom(&cpuIndex);
258 EXPECT_EQ(cpuIndex.ntotal, gpuIndex.ntotal);
259 EXPECT_EQ(gpuIndex.ntotal, opt.numAdd);
261 EXPECT_EQ(cpuIndex.d, gpuIndex.d);
262 EXPECT_EQ(cpuIndex.d, opt.dim);
263 EXPECT_EQ(cpuIndex.nlist, gpuIndex.getNumLists());
264 EXPECT_EQ(cpuIndex.nprobe, gpuIndex.getNumProbes());
265 EXPECT_EQ(cpuIndex.pq.M, gpuIndex.getNumSubQuantizers());
266 EXPECT_EQ(gpuIndex.getNumSubQuantizers(), opt.codes);
267 EXPECT_EQ(cpuIndex.pq.nbits, gpuIndex.getBitsPerCode());
268 EXPECT_EQ(gpuIndex.getBitsPerCode(), opt.bitsPerCode);
271 faiss::gpu::compareIndices(cpuIndex, gpuIndex,
272 opt.numQuery, opt.dim, opt.k, opt.toString(),
273 opt.getCompareEpsilon(),
274 opt.getPctMaxDiff1(),
275 opt.getPctMaxDiffN());
278 TEST(TestGpuIndexIVFPQ, QueryNaN) {
279 faiss::gpu::newTestSeed();
283 std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
284 std::vector<float> addVecs = faiss::gpu::randVecs(opt.numAdd, opt.dim);
299 gpuIndex.setPrecomputedCodes(opt.usePrecomputed);
300 gpuIndex.setNumProbes(opt.nprobe);
302 gpuIndex.train(opt.numTrain, trainVecs.data());
303 gpuIndex.add(opt.numAdd, addVecs.data());
306 std::vector<float> nans(numQuery * opt.dim,
307 std::numeric_limits<float>::quiet_NaN());
309 std::vector<float> distances(numQuery * opt.k, 0);
310 std::vector<faiss::Index::idx_t> indices(numQuery * opt.k, 0);
312 gpuIndex.search(numQuery,
318 for (
int q = 0; q < numQuery; ++q) {
319 for (
int k = 0; k < opt.k; ++k) {
320 EXPECT_EQ(indices[q * opt.k + k], -1);
321 EXPECT_EQ(distances[q * opt.k + k], std::numeric_limits<float>::max());
326 TEST(TestGpuIndexIVFPQ, AddNaN) {
327 faiss::gpu::newTestSeed();
344 gpuIndex.setPrecomputedCodes(opt.usePrecomputed);
345 gpuIndex.setNumProbes(opt.nprobe);
348 std::vector<float> nans(numNans * opt.dim,
349 std::numeric_limits<float>::quiet_NaN());
352 for (
int i = 0; i < opt.dim; ++i) {
356 std::vector<float> trainVecs = faiss::gpu::randVecs(opt.numTrain, opt.dim);
357 gpuIndex.train(opt.numTrain, trainVecs.data());
360 EXPECT_EQ(gpuIndex.ntotal, 0);
361 gpuIndex.add(numNans, nans.data());
364 EXPECT_EQ(gpuIndex.ntotal, 1);
366 std::vector<float> queryVecs = faiss::gpu::randVecs(opt.numQuery, opt.dim);
367 std::vector<float> distance(opt.numQuery * opt.k, 0);
368 std::vector<faiss::Index::idx_t> indices(opt.numQuery * opt.k, 0);
371 gpuIndex.search(opt.numQuery, queryVecs.data(), opt.k,
372 distance.data(), indices.data());