Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
PerfIVFPQAdd.cpp
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
10 #include <cuda_profiler_api.h>
11 #include "../../IndexFlat.h"
12 #include "../../IndexIVFPQ.h"
13 #include "../GpuIndexIVFPQ.h"
14 #include "../StandardGpuResources.h"
15 #include "../test/TestUtils.h"
16 #include "../utils/DeviceUtils.h"
17 #include "../utils/Timer.h"
18 #include <gflags/gflags.h>
19 #include <map>
20 #include <vector>
21
22 DEFINE_int32(batches, 10, "number of batches of vectors to add");
23 DEFINE_int32(batch_size, 10000, "number of vectors in each batch");
24 DEFINE_int32(dim, 256, "dimension of vectors");
25 DEFINE_int32(centroids, 4096, "num coarse centroids to use");
26 DEFINE_int32(bytes_per_vec, 32, "bytes per encoded vector");
27 DEFINE_int32(bits_per_code, 8, "bits per PQ code");
28 DEFINE_int32(index, 2, "0 = no indices on GPU; 1 = 32 bit, 2 = 64 bit on GPU");
29 DEFINE_bool(time_gpu, true, "time add to GPU");
30 DEFINE_bool(time_cpu, false, "time add to CPU");
31 DEFINE_bool(per_batch_time, false, "print per-batch times");
32 DEFINE_bool(reserve_memory, false, "whether or not to pre-reserve memory");
33
34 int main(int argc, char** argv) {
35 gflags::ParseCommandLineFlags(&argc, &argv, true);
36
37 cudaProfilerStop();
38
39 int dim = FLAGS_dim;
40 int numCentroids = FLAGS_centroids;
41 int bytesPerVec = FLAGS_bytes_per_vec;
42 int bitsPerCode = FLAGS_bits_per_code;
43
44 faiss::gpu::StandardGpuResources res;
45
46 // IndexIVFPQ will complain, but just give us enough to get through this
47 int numTrain = 4 * numCentroids;
48 std::vector<float> trainVecs = faiss::gpu::randVecs(numTrain, dim);
49
50 faiss::IndexFlatL2 coarseQuantizer(dim);
51 faiss::IndexIVFPQ cpuIndex(&coarseQuantizer, dim, numCentroids,
52 bytesPerVec, bitsPerCode);
53 if (FLAGS_time_cpu) {
54 cpuIndex.train(numTrain, trainVecs.data());
55 }
56
57 faiss::gpu::GpuIndexIVFPQConfig config;
58 config.device = 0;
59 config.indicesOptions = (faiss::gpu::IndicesOptions) FLAGS_index;
60
61 faiss::gpu::GpuIndexIVFPQ gpuIndex(
62 &res, dim, numCentroids, bytesPerVec, bitsPerCode,
63 faiss::METRIC_L2, config);
64
65 if (FLAGS_time_gpu) {
66 gpuIndex.train(numTrain, trainVecs.data());
67 if (FLAGS_reserve_memory) {
68 size_t numVecs = (size_t) FLAGS_batches * (size_t) FLAGS_batch_size;
69 gpuIndex.reserveMemory(numVecs);
70 }
71 }
72
73 cudaDeviceSynchronize();
74 CUDA_VERIFY(cudaProfilerStart());
75
76 float totalGpuTime = 0.0f;
77 float totalCpuTime = 0.0f;
78
79 for (int i = 0; i < FLAGS_batches; ++i) {
80 if (!FLAGS_per_batch_time) {
81 if (i % 10 == 0) {
82 printf("Adding batch %d\n", i + 1);
83 }
84 }
85
86 auto addVecs = faiss::gpu::randVecs(FLAGS_batch_size, dim);
87
88 if (FLAGS_time_gpu) {
89 faiss::gpu::CpuTimer timer;
90 gpuIndex.add(FLAGS_batch_size, addVecs.data());
91 CUDA_VERIFY(cudaDeviceSynchronize());
92 auto time = timer.elapsedMilliseconds();
93
94 totalGpuTime += time;
95
96 if (FLAGS_per_batch_time) {
97 printf("Batch %d | GPU time to add %d vecs: %.3f ms (%.5f ms per)\n",
98 i + 1, FLAGS_batch_size, time, time / (float) FLAGS_batch_size);
99 }
100 }
101
102 if (FLAGS_time_cpu) {
103 faiss::gpu::CpuTimer timer;
104 cpuIndex.add(FLAGS_batch_size, addVecs.data());
105 auto time = timer.elapsedMilliseconds();
106
107 totalCpuTime += time;
108
109 if (FLAGS_per_batch_time) {
110 printf("Batch %d | CPU time to add %d vecs: %.3f ms (%.5f ms per)\n",
111 i + 1, FLAGS_batch_size, time, time / (float) FLAGS_batch_size);
112 }
113 }
114 }
115
116 CUDA_VERIFY(cudaProfilerStop());
117
118 int total = FLAGS_batch_size * FLAGS_batches;
119
120 if (FLAGS_time_gpu) {
121 printf("%d dim, %d centroids, %d x %d encoding\n"
122 "GPU time to add %d vectors (%d batches, %d per batch): "
123 "%.3f ms (%.3f us per)\n",
124 dim, numCentroids, bytesPerVec, bitsPerCode,
125 total, FLAGS_batches, FLAGS_batch_size,
126 totalGpuTime, totalGpuTime * 1000.0f / (float) total);
127 }
128
129 if (FLAGS_time_cpu) {
130 printf("%d dim, %d centroids, %d x %d encoding\n"
131 "CPU time to add %d vectors (%d batches, %d per batch): "
132 "%.3f ms (%.3f us per)\n",
133 dim, numCentroids, bytesPerVec, bitsPerCode,
134 total, FLAGS_batches, FLAGS_batch_size,
135 totalCpuTime, totalCpuTime * 1000.0f / (float) total);
136 }
137
138 return 0;
139 }
faiss::gpu::CpuTimer::elapsedMilliseconds
float elapsedMilliseconds()
Returns elapsed time in milliseconds.
Definition: Timer.cpp:50
faiss::IndexFlatL2
Definition: IndexFlat.h:77
faiss::gpu::CpuTimer
CPU wallclock elapsed timer.
Definition: Timer.h:40
faiss::gpu::GpuIndexConfig::device
int device
GPU device on which the index is resident.
Definition: GpuIndex.h:25
faiss::gpu::GpuIndexIVFPQConfig
Definition: GpuIndexIVFPQ.h:21
faiss::IndexIVFPQ
Definition: IndexIVFPQ.h:34
faiss::gpu::StandardGpuResources
Definition: StandardGpuResources.h:21
faiss::gpu::GpuIndexIVFPQ
IVFPQ index for the GPU.
Definition: GpuIndexIVFPQ.h:38
faiss::gpu::GpuIndexIVFConfig::indicesOptions
IndicesOptions indicesOptions
Index storage options for the GPU.
Definition: GpuIndexIVF.h:29