Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
TestGpuIndexFlat.cpp
1 /**
2 * Copyright (c) 2015-present, Facebook, Inc.
3 * All rights reserved.
4 *
5 * This source code is licensed under the CC-by-NC license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9 // Copyright 2004-present Facebook. All Rights Reserved.
10
11 #include "../../IndexFlat.h"
12 #include "../GpuIndexFlat.h"
13 #include "../StandardGpuResources.h"
14 #include "../utils/DeviceUtils.h"
15 #include "../test/TestUtils.h"
16 #include <gtest/gtest.h>
17 #include <sstream>
18 #include <vector>
19
20 // FIXME: figure out a better way to test fp16
21 constexpr float kF16MaxRelErr = 0.07f;
22 constexpr float kF32MaxRelErr = 6e-3f;
23
24 void testFlat(bool useL2,
25 bool useFloat16,
26 bool useTransposed,
27 int kOverride = -1) {
28 int numVecs = faiss::gpu::randVal(1000, 20000);
29 int dim = faiss::gpu::randVal(50, 800);
30 int numQuery = faiss::gpu::randVal(1, 512);
31
32 // Due to loss of precision in a float16 accumulator, for large k,
33 // the number of differences is pretty huge. Restrict ourselves to a
34 // fairly small `k` for float16
35 int k = useFloat16 ?
36 std::min(faiss::gpu::randVal(1, 50), numVecs) :
37 std::min(faiss::gpu::randVal(1, 1024), numVecs);
38 if (kOverride > 0) {
39 k = kOverride;
40 }
41
42 faiss::IndexFlatIP cpuIndexIP(dim);
43 faiss::IndexFlatL2 cpuIndexL2(dim);
44
45 faiss::IndexFlat* cpuIndex =
46 useL2 ? (faiss::IndexFlat*) &cpuIndexL2 : (faiss::IndexFlat*) &cpuIndexIP;
47
48 // Construct on a random device to test multi-device, if we have
49 // multiple devices
50 int device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
51
52 faiss::gpu::StandardGpuResources res;
53 res.noTempMemory();
54
55
56 faiss::gpu::GpuIndexFlatConfig config;
57 config.device = device;
58 config.useFloat16 = useFloat16;
59 config.storeTransposed = useTransposed;
60
61 faiss::gpu::GpuIndexFlatIP gpuIndexIP(&res, dim, config);
62 faiss::gpu::GpuIndexFlatL2 gpuIndexL2(&res, dim, config);
63
64 faiss::gpu::GpuIndexFlat* gpuIndex =
65 useL2 ? (faiss::gpu::GpuIndexFlat*) &gpuIndexL2 :
66 (faiss::gpu::GpuIndexFlat*) &gpuIndexIP;
67
68 std::vector<float> vecs = faiss::gpu::randVecs(numVecs, dim);
69 cpuIndex->add(numVecs, vecs.data());
70 gpuIndex->add(numVecs, vecs.data());
71
72 std::stringstream str;
73 str << (useL2 ? "L2" : "IP") << " numVecs " << numVecs
74 << " dim " << dim
75 << " useFloat16 " << useFloat16
76 << " transposed " << useTransposed
77 << " numQuery " << numQuery
78 << " k " << k;
79
80 // To some extent, we depend upon the relative error for the test
81 // for float16
82 faiss::gpu::compareIndices(*cpuIndex, *gpuIndex, numQuery, dim, k, str.str(),
83 useFloat16 ? kF16MaxRelErr : kF32MaxRelErr,
84 // FIXME: the fp16 bounds are
85 // useless when math (the accumulator) is
86 // in fp16. Figure out another way to test
87 useFloat16 ? 0.99f : 0.1f,
88 useFloat16 ? 0.65f : 0.015f);
89 }
90
91 TEST(TestGpuIndexFlat, IP_Float32) {
92 for (int tries = 0; tries < 5; ++tries) {
93 faiss::gpu::newTestSeed();
94 testFlat(false, false, false);
95 testFlat(false, false, true);
96 }
97 }
98
99 TEST(TestGpuIndexFlat, L2_Float32) {
100 for (int tries = 0; tries < 5; ++tries) {
101 faiss::gpu::newTestSeed();
102 testFlat(true, false, false);
103 testFlat(true, false, true);
104 }
105 }
106
107 // test specialized k == 1 codepath
108 TEST(TestGpuIndexFlat, L2_Float32_K1) {
109 for (int tries = 0; tries < 5; ++tries) {
110 faiss::gpu::newTestSeed();
111 testFlat(true, false, false, 1);
112 testFlat(true, false, true, 1);
113 }
114 }
115
116 TEST(TestGpuIndexFlat, IP_Float16) {
117 for (int tries = 0; tries < 5; ++tries) {
118 faiss::gpu::newTestSeed();
119 testFlat(false, true, false);
120 testFlat(false, true, false);
121 }
122 }
123
124 TEST(TestGpuIndexFlat, L2_Float16) {
125 for (int tries = 0; tries < 5; ++tries) {
126 faiss::gpu::newTestSeed();
127 testFlat(true, true, false);
128 testFlat(true, true, true);
129 }
130 }
131
132 // test specialized k == 1 codepath
133 TEST(TestGpuIndexFlat, L2_Float16_K1) {
134 for (int tries = 0; tries < 5; ++tries) {
135 faiss::gpu::newTestSeed();
136 testFlat(true, true, false, 1);
137 testFlat(true, true, true, 1);
138 }
139 }
140
141 TEST(TestGpuIndexFlat, QueryEmpty) {
142 faiss::gpu::StandardGpuResources res;
143 res.noTempMemory();
144
145 faiss::gpu::GpuIndexFlatConfig config;
146 config.device = 0;
147 config.useFloat16 = false;
148 config.storeTransposed = false;
149
150 int dim = 128;
151 faiss::gpu::GpuIndexFlatL2 gpuIndex(&res, dim, config);
152
153 // Querying an empty index should not blow up, and just return
154 // (FLT_MAX, -1)
155 int numQuery = 10;
156 int k = 50;
157 std::vector<float> queries(numQuery * dim, 1.0f);
158
159 std::vector<float> dist(numQuery * k, 0);
160 std::vector<faiss::Index::idx_t> ind(numQuery * k);
161
162 gpuIndex.search(numQuery, queries.data(), k, dist.data(), ind.data());
163
164 for (auto d : dist) {
165 EXPECT_EQ(d, std::numeric_limits<float>::max());
166 }
167
168 for (auto i : ind) {
169 EXPECT_EQ(i, -1);
170 }
171 }
172
173 TEST(TestGpuIndexFlat, CopyFrom) {
174 faiss::gpu::newTestSeed();
175
176 int numVecs = faiss::gpu::randVal(100, 200);
177 int dim = faiss::gpu::randVal(1, 1000);
178
179 faiss::IndexFlatL2 cpuIndex(dim);
180
181 std::vector<float> vecs = faiss::gpu::randVecs(numVecs, dim);
182 cpuIndex.add(numVecs, vecs.data());
183
184 faiss::gpu::StandardGpuResources res;
185 res.noTempMemory();
186
187 // Fill with garbage values
188 int device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
189
190 faiss::gpu::GpuIndexFlatConfig config;
191 config.device = 0;
192 config.useFloat16 = false;
193 config.storeTransposed = false;
194
195 faiss::gpu::GpuIndexFlatL2 gpuIndex(&res, 2000, config);
196 gpuIndex.copyFrom(&cpuIndex);
197
198 EXPECT_EQ(cpuIndex.ntotal, gpuIndex.ntotal);
199 EXPECT_EQ(gpuIndex.ntotal, numVecs);
200
201 EXPECT_EQ(cpuIndex.d, gpuIndex.d);
202 EXPECT_EQ(cpuIndex.d, dim);
203
204 int idx = faiss::gpu::randVal(0, numVecs - 1);
205
206 std::vector<float> gpuVals(dim);
207 gpuIndex.reconstruct(idx, gpuVals.data());
208
209 std::vector<float> cpuVals(dim);
210 cpuIndex.reconstruct(idx, cpuVals.data());
211
212 EXPECT_EQ(gpuVals, cpuVals);
213 }
214
215 TEST(TestGpuIndexFlat, CopyTo) {
216 faiss::gpu::newTestSeed();
217
218 faiss::gpu::StandardGpuResources res;
219 res.noTempMemory();
220
221 int numVecs = faiss::gpu::randVal(100, 200);
222 int dim = faiss::gpu::randVal(1, 1000);
223
224 int device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
225
226 faiss::gpu::GpuIndexFlatConfig config;
227 config.device = device;
228 config.useFloat16 = false;
229 config.storeTransposed = false;
230
231 faiss::gpu::GpuIndexFlatL2 gpuIndex(&res, dim, config);
232
233 std::vector<float> vecs = faiss::gpu::randVecs(numVecs, dim);
234 gpuIndex.add(numVecs, vecs.data());
235
236 // Fill with garbage values
237 faiss::IndexFlatL2 cpuIndex(2000);
238 gpuIndex.copyTo(&cpuIndex);
239
240 EXPECT_EQ(cpuIndex.ntotal, gpuIndex.ntotal);
241 EXPECT_EQ(gpuIndex.ntotal, numVecs);
242
243 EXPECT_EQ(cpuIndex.d, gpuIndex.d);
244 EXPECT_EQ(cpuIndex.d, dim);
245
246 int idx = faiss::gpu::randVal(0, numVecs - 1);
247
248 std::vector<float> gpuVals(dim);
249 gpuIndex.reconstruct(idx, gpuVals.data());
250
251 std::vector<float> cpuVals(dim);
252 cpuIndex.reconstruct(idx, cpuVals.data());
253
254 EXPECT_EQ(gpuVals, cpuVals);
255 }
256
257 TEST(TestGpuIndexFlat, UnifiedMemory) {
258 // Construct on a random device to test multi-device, if we have
259 // multiple devices
260 int device = faiss::gpu::randVal(0, faiss::gpu::getNumDevices() - 1);
261
262 if (!faiss::gpu::getFullUnifiedMemSupport(device)) {
263 return;
264 }
265
266 int dim = 256;
267
268 // FIXME: GpuIndexFlat doesn't support > 2^31 (vecs * dims) due to
269 // kernel indexing, so we can't test unified memory for memory
270 // oversubscription.
271 size_t numVecs = 50000;
272 int numQuery = 10;
273 int k = 10;
274
275 faiss::IndexFlatL2 cpuIndexL2(dim);
276
277 faiss::gpu::StandardGpuResources res;
278 res.noTempMemory();
279
280 faiss::gpu::GpuIndexFlatConfig config;
281 config.device = device;
282 config.memorySpace = faiss::gpu::MemorySpace::Unified;
283
284 faiss::gpu::GpuIndexFlatL2 gpuIndexL2(&res, dim, config);
285
286 std::vector<float> vecs = faiss::gpu::randVecs(numVecs, dim);
287 cpuIndexL2.add(numVecs, vecs.data());
288 gpuIndexL2.add(numVecs, vecs.data());
289
290 // To some extent, we depend upon the relative error for the test
291 // for float16
292 faiss::gpu::compareIndices(cpuIndexL2, gpuIndexL2,
293 numQuery, dim, k, "Unified Memory",
294 kF32MaxRelErr,
295 0.1f,
296 0.015f);
297 }
faiss::IndexFlat
Definition: IndexFlat.h:23
faiss::gpu::GpuIndexFlatConfig::storeTransposed
bool storeTransposed
Definition: GpuIndexFlat.h:48
faiss::gpu::GpuIndexFlat::copyTo
void copyTo(faiss::IndexFlat *index) const
Definition: GpuIndexFlat.cu:120
faiss::gpu::StandardGpuResources::noTempMemory
void noTempMemory()
Definition: StandardGpuResources.cpp:71
faiss::IndexFlatL2
Definition: IndexFlat.h:79
faiss::IndexFlat::reconstruct
void reconstruct(idx_t key, float *recons) const override
Definition: IndexFlat.cpp:118
faiss::IndexFlatIP
Definition: IndexFlat.h:73
faiss::gpu::GpuIndexFlatConfig::useFloat16
bool useFloat16
Whether or not data is stored as float16.
Definition: GpuIndexFlat.h:35
faiss::gpu::GpuIndexConfig::device
int device
GPU device on which the index is resident.
Definition: GpuIndex.h:27
faiss::Index::d
int d
vector dimension
Definition: Index.h:64
faiss::gpu::GpuIndexFlatConfig
Definition: GpuIndexFlat.h:27
faiss::gpu::GpuIndexConfig::memorySpace
MemorySpace memorySpace
Definition: GpuIndex.h:32
faiss::gpu::StandardGpuResources
Definition: StandardGpuResources.h:23
faiss::gpu::GpuIndexFlat::reconstruct
void reconstruct(faiss::Index::idx_t key, float *out) const override
Definition: GpuIndexFlat.cu:491
faiss::Index::ntotal
idx_t ntotal
total nb of indexed vectors
Definition: Index.h:65
faiss::IndexFlat::add
void add(idx_t n, const float *x) override
Definition: IndexFlat.cpp:30
faiss::gpu::GpuIndexFlatIP
Definition: GpuIndexFlat.h:194
faiss::gpu::GpuIndexFlat::copyFrom
void copyFrom(const faiss::IndexFlat *index)
Definition: GpuIndexFlat.cu:87
faiss::gpu::GpuIndexFlatL2
Definition: GpuIndexFlat.h:169
faiss::gpu::GpuIndexFlat::add
void add(faiss::Index::idx_t, const float *x) override
Overrides to avoid excessive copies.
Definition: GpuIndexFlat.cu:162
faiss::gpu::GpuIndexFlat
Definition: GpuIndexFlat.h:54
faiss::gpu::GpuIndexFlat::search
void search(faiss::Index::idx_t n, const float *x, faiss::Index::idx_t k, float *distances, faiss::Index::idx_t *labels) const override
Definition: GpuIndexFlat.cu:205