Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

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