Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
PerfFlat.cu
1 /**
2 * Copyright (c) Facebook, Inc. and its affiliates.
3 *
4 * This source code is licensed under the MIT license found in the
5 * LICENSE file in the root directory of this source tree.
6 */
7
8
9 #include "../../IndexFlat.h"
10 #include "../../utils.h"
11 #include "../GpuIndexFlat.h"
12 #include "IndexWrapper.h"
13 #include "../test/TestUtils.h"
14 #include "../utils/DeviceTensor.cuh"
15 #include "../utils/DeviceUtils.h"
16 #include "../utils/HostTensor.cuh"
17 #include "../utils/Timer.h"
18 #include <gflags/gflags.h>
19 #include <map>
20 #include <memory>
21 #include <vector>
22
23 #include <cuda_profiler_api.h>
24
25 DEFINE_bool(l2, true, "L2 or inner product");
26 DEFINE_int32(k, 3, "final number of closest results returned");
27 DEFINE_int32(num, 128, "# of vecs");
28 DEFINE_int32(dim, 128, "# of dimensions");
29 DEFINE_int32(num_queries, 3, "number of query vectors");
30 DEFINE_bool(diff, true, "show exact distance + index output discrepancies");
31 DEFINE_bool(use_float16, false, "use encodings in float16");
32 DEFINE_bool(use_float16_math, false, "perform math in float16");
33 DEFINE_bool(transposed, false, "store vectors transposed");
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 DEFINE_bool(cpu, true, "run the CPU code for timing and comparison");
38 DEFINE_bool(use_unified_mem, false, "use Pascal unified memory for the index");
39
40 using namespace faiss::gpu;
41
42 int main(int argc, char** argv) {
43 gflags::ParseCommandLineFlags(&argc, &argv, true);
44
45 cudaProfilerStop();
46
47 auto seed = FLAGS_seed != -1L ? FLAGS_seed : time(nullptr);
48 printf("using seed %ld\n", seed);
49
50 auto numQueries = FLAGS_num_queries;
51
52 auto index = std::unique_ptr<faiss::IndexFlat>(
53 new faiss::IndexFlat(FLAGS_dim, FLAGS_l2 ?
54 faiss::METRIC_L2 : faiss::METRIC_INNER_PRODUCT));
55
56 HostTensor<float, 2, true> vecs({FLAGS_num, FLAGS_dim});
57 faiss::float_rand(vecs.data(), vecs.numElements(), seed);
58
59 index->add(FLAGS_num, vecs.data());
60
61 printf("Database: dim %d num vecs %d\n", FLAGS_dim, FLAGS_num);
62 printf("%s lookup: %d queries, total k %d\n",
63 FLAGS_l2 ? "L2" : "IP",
64 numQueries, FLAGS_k);
65 printf("float16 encoding %s\n", FLAGS_use_float16 ? "enabled" : "disabled");
66 printf("transposed storage %s\n", FLAGS_transposed ? "enabled" : "disabled");
67
68 // Convert to GPU index
69 printf("Copying index to %d GPU(s)...\n", FLAGS_num_gpus);
70
71 auto initFn = [&index](faiss::gpu::GpuResources* res, int dev) ->
72 std::unique_ptr<faiss::gpu::GpuIndexFlat> {
73 ((faiss::gpu::StandardGpuResources*) res)->setPinnedMemory(
74 FLAGS_pinned_mem);
75
76 GpuIndexFlatConfig config;
77 config.device = dev;
78 config.useFloat16 = FLAGS_use_float16;
79 config.useFloat16Accumulator = FLAGS_use_float16_math;
80 config.storeTransposed = FLAGS_transposed;
81 config.memorySpace = FLAGS_use_unified_mem ?
82 MemorySpace::Unified : MemorySpace::Device;
83
84 auto p = std::unique_ptr<faiss::gpu::GpuIndexFlat>(
85 new faiss::gpu::GpuIndexFlat(res, index.get(), config));
86 return p;
87 };
88
89 IndexWrapper<faiss::gpu::GpuIndexFlat> gpuIndex(FLAGS_num_gpus, initFn);
90 printf("copy done\n");
91
92 // Build query vectors
93 HostTensor<float, 2, true> cpuQuery({numQueries, FLAGS_dim});
94 faiss::float_rand(cpuQuery.data(), cpuQuery.numElements(), seed);
95
96 // Time faiss CPU
97 HostTensor<float, 2, true> cpuDistances({numQueries, FLAGS_k});
98 HostTensor<faiss::Index::idx_t, 2, true> cpuIndices({numQueries, FLAGS_k});
99
100 if (FLAGS_cpu) {
101 float cpuTime = 0.0f;
102
103 CpuTimer timer;
104 index->search(numQueries,
105 cpuQuery.data(),
106 FLAGS_k,
107 cpuDistances.data(),
108 cpuIndices.data());
109
110 cpuTime = timer.elapsedMilliseconds();
111 printf("CPU time %.3f ms\n", cpuTime);
112 }
113
114 HostTensor<float, 2, true> gpuDistances({numQueries, FLAGS_k});
115 HostTensor<faiss::Index::idx_t, 2, true> gpuIndices({numQueries, FLAGS_k});
116
117 CUDA_VERIFY(cudaProfilerStart());
118 faiss::gpu::synchronizeAllDevices();
119
120 float gpuTime = 0.0f;
121
122 // Time GPU
123 {
124 CpuTimer timer;
125
126 gpuIndex.getIndex()->search(cpuQuery.getSize(0),
127 cpuQuery.data(),
128 FLAGS_k,
129 gpuDistances.data(),
130 gpuIndices.data());
131
132 // There is a device -> host copy above, so no need to time
133 // additional synchronization with the GPU
134 gpuTime = timer.elapsedMilliseconds();
135 }
136
137 CUDA_VERIFY(cudaProfilerStop());
138 printf("GPU time %.3f ms\n", gpuTime);
139
140 if (FLAGS_cpu) {
141 compareLists(cpuDistances.data(), cpuIndices.data(),
142 gpuDistances.data(), gpuIndices.data(),
143 numQueries, FLAGS_k,
144 "", true, FLAGS_diff, false);
145 }
146
147 CUDA_VERIFY(cudaDeviceSynchronize());
148
149 return 0;
150 }
faiss::gpu::CpuTimer::elapsedMilliseconds
float elapsedMilliseconds()
Returns elapsed time in milliseconds.
Definition: Timer.cpp:50
faiss::IndexFlat
Definition: IndexFlat.h:21
faiss::gpu::GpuIndexFlatConfig::useFloat16Accumulator
bool useFloat16Accumulator
Definition: GpuIndexFlat.h:38
faiss::gpu::GpuIndexFlatConfig::storeTransposed
bool storeTransposed
Definition: GpuIndexFlat.h:46
faiss::gpu::GpuResources
Definition: GpuResources.h:21
faiss::gpu::CpuTimer
CPU wallclock elapsed timer.
Definition: Timer.h:40
faiss::gpu::GpuIndexFlatConfig::useFloat16
bool useFloat16
Whether or not data is stored as float16.
Definition: GpuIndexFlat.h:33
faiss::gpu::GpuIndexConfig::device
int device
GPU device on which the index is resident.
Definition: GpuIndex.h:25
faiss::gpu::GpuIndexFlatConfig
Definition: GpuIndexFlat.h:25
faiss::gpu::GpuIndexConfig::memorySpace
MemorySpace memorySpace
Definition: GpuIndex.h:30
faiss::gpu::StandardGpuResources
Definition: StandardGpuResources.h:21
faiss::gpu::HostTensor
Definition: HostTensor.cuh:20
faiss::gpu::GpuIndexFlat
Definition: GpuIndexFlat.h:52
faiss::gpu::IndexWrapper
Definition: IndexWrapper.h:23