Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
TestGpuIndexBinaryFlat.cpp
1 /**
2 * Copyright (c) 2015-present, Facebook, Inc.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD+Patents license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9
10 #include "../../IndexBinaryFlat.h"
11 #include "../GpuIndexBinaryFlat.h"
12 #include "../StandardGpuResources.h"
13 #include "../utils/DeviceUtils.h"
14 #include "../test/TestUtils.h"
15 #include "../../utils.h"
16 #include <gtest/gtest.h>
17 #include <sstream>
18 #include <vector>
19
20 void compareBinaryDist(const std::vector<int>& cpuDist,
21 const std::vector<faiss::IndexBinary::idx_t>& cpuLabels,
22 const std::vector<int>& gpuDist,
23 const std::vector<faiss::IndexBinary::idx_t>& gpuLabels,
24 int numQuery,
25 int k) {
26 for (int i = 0; i < numQuery; ++i) {
27 // The index order can be permuted within a group that has the same
28 // distance, since this is based on the order in which the algorithm
29 // encounters the values. The last set of equivalent distances seen in the
30 // min-k might be truncated, so we can't check that set, but all others we
31 // can check.
32 std::set<faiss::IndexBinary::idx_t> cpuLabelSet;
33 std::set<faiss::IndexBinary::idx_t> gpuLabelSet;
34
35 int curDist = -1;
36
37 for (int j = 0; j < k; ++j) {
38 int idx = i * k + j;
39
40 if (curDist == -1) {
41 curDist = cpuDist[idx];
42 }
43
44 if (curDist != cpuDist[idx]) {
45 // Distances must be monotonically increasing
46 EXPECT_LT(curDist, cpuDist[idx]);
47
48 // This is a new set of distances
49 EXPECT_EQ(cpuLabelSet, gpuLabelSet);
50 curDist = cpuDist[idx];
51 cpuLabelSet.clear();
52 gpuLabelSet.clear();
53 }
54
55 cpuLabelSet.insert(cpuLabels[idx]);
56 gpuLabelSet.insert(gpuLabels[idx]);
57
58 // Because the distances are reproducible, they must be exactly the same
59 EXPECT_EQ(cpuDist[idx], gpuDist[idx]);
60 }
61 }
62 }
63
64 template <int DimMultiple>
65 void testGpuIndexBinaryFlat() {
66 faiss::gpu::StandardGpuResources res;
67 res.noTempMemory();
68
69 faiss::gpu::GpuIndexBinaryFlatConfig config;
70 config.device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
71
72 // multiples of 8 and multiples of 32 use different implementations
73 int dims = faiss::gpu::randVal(1, 20) * DimMultiple;
74 faiss::gpu::GpuIndexBinaryFlat gpuIndex(&res, dims, config);
75
76 faiss::IndexBinaryFlat cpuIndex(dims);
77
78 int numVecs = faiss::gpu::randVal(1, 20000);
79 int numQuery = faiss::gpu::randVal(1, 1000);
80 int k = faiss::gpu::randVal(1, 1024);
81
82 auto data = faiss::gpu::randBinaryVecs(numVecs, dims);
83 gpuIndex.add(numVecs, data.data());
84 cpuIndex.add(numVecs, data.data());
85
86 auto query = faiss::gpu::randBinaryVecs(numQuery, dims);
87
88 std::vector<int> cpuDist(numQuery * k);
89 std::vector<faiss::IndexBinary::idx_t> cpuLabels(numQuery * k);
90
91 cpuIndex.search(numQuery,
92 query.data(),
93 k,
94 cpuDist.data(),
95 cpuLabels.data());
96
97 std::vector<int> gpuDist(numQuery * k);
98 std::vector<faiss::IndexBinary::idx_t> gpuLabels(numQuery * k);
99
100 gpuIndex.search(numQuery,
101 query.data(),
102 k,
103 gpuDist.data(),
104 gpuLabels.data());
105
106 compareBinaryDist(cpuDist, cpuLabels,
107 gpuDist, gpuLabels,
108 numQuery, k);
109 }
110
111 TEST(TestGpuIndexBinaryFlat, Test8) {
112 for (int tries = 0; tries < 4; ++tries) {
113 testGpuIndexBinaryFlat<8>();
114 }
115 }
116
117 TEST(TestGpuIndexBinaryFlat, Test32) {
118 for (int tries = 0; tries < 4; ++tries) {
119 testGpuIndexBinaryFlat<32>();
120 }
121 }
122
123 int main(int argc, char** argv) {
124 testing::InitGoogleTest(&argc, argv);
125
126 // just run with a fixed test seed
127 faiss::gpu::setTestSeed(100);
128
129 return RUN_ALL_TESTS();
130 }
faiss::IndexBinaryFlat
Definition: IndexBinaryFlat.h:22
faiss::gpu::GpuIndexBinaryFlatConfig
Definition: GpuIndexBinaryFlat.h:19
faiss::gpu::StandardGpuResources::noTempMemory
void noTempMemory()
Definition: StandardGpuResources.cpp:76
faiss::gpu::GpuIndexConfig::device
int device
GPU device on which the index is resident.
Definition: GpuIndex.h:26
faiss::gpu::StandardGpuResources
Definition: StandardGpuResources.h:22
faiss::gpu::GpuIndexBinaryFlat
Definition: GpuIndexBinaryFlat.h:24