Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
CompareFlat.cu
1
2 /**
3 * Copyright (c) 2015-present, Facebook, Inc.
4 * All rights reserved.
5 *
6 * This source code is licensed under the CC-by-NC license found in the
7 * LICENSE file in the root directory of this source tree.
8 */
9
10 // Copyright 2004-present Facebook. All Rights Reserved.
11
12 #include "../../IndexFlat.h"
13 #include "../../utils.h"
14 #include "../GpuIndexFlat.h"
15 #include "IndexWrapper.h"
16 #include "../test/TestUtils.h"
17 #include "../utils/DeviceTensor.cuh"
18 #include "../utils/DeviceUtils.h"
19 #include "../utils/HostTensor.cuh"
20 #include "../utils/Timer.h"
21 #include <gflags/gflags.h>
22 #include <map>
23 #include <memory>
24 #include <vector>
25
26 #include <cuda_profiler_api.h>
27
28 DEFINE_int32(k, 3, "final number of closest results returned");
29 DEFINE_int32(num, 128, "# of vecs");
30 DEFINE_int32(dim, 128, "# of dimensions");
31 DEFINE_int32(num_queries, 3, "number of query vectors");
32 DEFINE_bool(diff, true, "show exact distance + index output discrepancies");
33 DEFINE_bool(use_float16, false, "use encodings in float16 instead of float32");
34 DEFINE_int64(seed, -1, "specify random seed");
35 DEFINE_int32(num_gpus, 1, "number of gpus to use");
36 DEFINE_int64(pinned_mem, 0, "pinned memory allocation to use");
37
38 using namespace faiss::gpu;
39
40 int main(int argc, char** argv) {
41 google::ParseCommandLineFlags(&argc, &argv, true);
42
43 cudaProfilerStop();
44
45 auto seed = FLAGS_seed != -1L ? FLAGS_seed : time(nullptr);
46 printf("using seed %ld\n", seed);
47
48 auto numQueries = FLAGS_num_queries;
49
50 auto index = std::unique_ptr<faiss::IndexFlatL2>(
51 new faiss::IndexFlatL2(FLAGS_dim));
52
53 HostTensor<float, 2, true> vecs({FLAGS_num, FLAGS_dim});
54 faiss::float_rand(vecs.data(), vecs.numElements(), seed);
55
56 index->add(FLAGS_num, vecs.data());
57
58 printf("Database: dim %d num vecs %d\n", FLAGS_dim, FLAGS_num);
59 printf("L2 lookup: %d queries, total k %d\n",
60 numQueries, FLAGS_k);
61 printf("float16 encoding %s\n", FLAGS_use_float16 ? "enabled" : "disabled");
62
63 // Convert to GPU index
64 printf("Copying index to %d GPU(s)...\n", FLAGS_num_gpus);
65
66 auto initFn = [&index](faiss::gpu::GpuResources* res, int dev) ->
67 std::unique_ptr<faiss::gpu::GpuIndexFlatL2> {
68 ((faiss::gpu::StandardGpuResources*) res)->setPinnedMemory(
69 FLAGS_pinned_mem);
70
71 auto p = std::unique_ptr<faiss::gpu::GpuIndexFlatL2>(
72 new faiss::gpu::GpuIndexFlatL2(res, dev, FLAGS_use_float16, index.get()));
73 return p;
74 };
75
76 IndexWrapper<faiss::gpu::GpuIndexFlatL2> gpuIndex(FLAGS_num_gpus, initFn);
77 printf("copy done\n");
78
79 // Build query vectors
80 HostTensor<float, 2, true> cpuQuery({numQueries, FLAGS_dim});
81 faiss::float_rand(cpuQuery.data(), cpuQuery.numElements(), seed);
82
83 // Time faiss CPU
84 HostTensor<float, 2, true> cpuDistances({numQueries, FLAGS_k});
85 HostTensor<faiss::Index::idx_t, 2, true> cpuIndices({numQueries, FLAGS_k});
86
87 float cpuTime = 0.0f;
88
89 {
90 CpuTimer timer;
91 index->search(numQueries,
92 cpuQuery.data(),
93 FLAGS_k,
94 cpuDistances.data(),
95 cpuIndices.data());
96
97 cpuTime = timer.elapsedMilliseconds();
98 }
99
100 printf("CPU time %.3f ms\n", cpuTime);
101
102 HostTensor<float, 2, true> gpuDistances({numQueries, FLAGS_k});
103 HostTensor<faiss::Index::idx_t, 2, true> gpuIndices({numQueries, FLAGS_k});
104
105 CUDA_VERIFY(cudaProfilerStart());
106 faiss::gpu::synchronizeAllDevices();
107
108 float gpuTime = 0.0f;
109
110 // Time GPU
111 {
112 CpuTimer timer;
113
114 gpuIndex.getIndex()->search(cpuQuery.getSize(0),
115 cpuQuery.data(),
116 FLAGS_k,
117 gpuDistances.data(),
118 gpuIndices.data());
119
120 // There is a device -> host copy above, so no need to time
121 // additional synchronization with the GPU
122 gpuTime = timer.elapsedMilliseconds();
123 }
124
125 CUDA_VERIFY(cudaProfilerStop());
126 printf("GPU time %.3f ms\n", gpuTime);
127
128 compareLists(cpuDistances.data(), cpuIndices.data(),
129 gpuDistances.data(), gpuIndices.data(),
130 numQueries, FLAGS_k,
131 "", true, FLAGS_diff, false);
132
133 CUDA_VERIFY(cudaDeviceSynchronize());
134 // printf("\ncudaMalloc usage %zd\n",
135 // resources.getMemoryManager().getHighWaterCudaMalloc());
136
137 return 0;
138 }
faiss::gpu::CpuTimer::elapsedMilliseconds
float elapsedMilliseconds()
Returns elapsed time in milliseconds.
Definition: Timer.cpp:53
faiss::IndexFlatL2
Definition: IndexFlat.h:81
faiss::gpu::GpuResources
Definition: GpuResources.h:24
faiss::gpu::CpuTimer
CPU wallclock elapsed timer.
Definition: Timer.h:43
faiss::gpu::StandardGpuResources
Definition: StandardGpuResources.h:24
faiss::gpu::GpuIndexFlatL2
Definition: GpuIndexFlat.h:131
faiss::gpu::HostTensor
Definition: HostTensor.cuh:23
faiss::gpu::IndexWrapper
Definition: IndexWrapper.h:26