Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
GpuIndexFlat.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 "GpuIndexFlat.h"
10 #include "../IndexFlat.h"
11 #include "GpuResources.h"
12 #include "impl/FlatIndex.cuh"
13 #include "utils/ConversionOperators.cuh"
14 #include "utils/CopyUtils.cuh"
15 #include "utils/DeviceUtils.h"
16 #include "utils/Float16.cuh"
17 #include "utils/StaticUtils.h"
18
19 #include <thrust/execution_policy.h>
20 #include <thrust/transform.h>
21 #include <limits>
22
23 namespace faiss { namespace gpu {
24
25 GpuIndexFlat::GpuIndexFlat(GpuResources* resources,
26 const faiss::IndexFlat* index,
27 GpuIndexFlatConfig config) :
28 GpuIndex(resources, index->d, index->metric_type, config),
29 config_(std::move(config)),
30 data_(nullptr) {
31 verifySettings_();
32
33 // Flat index doesn't need training
34 this->is_trained = true;
35
36 copyFrom(index);
37 }
38
39 GpuIndexFlat::GpuIndexFlat(GpuResources* resources,
40 int dims,
41 faiss::MetricType metric,
42 GpuIndexFlatConfig config) :
43 GpuIndex(resources, dims, metric, config),
44 config_(std::move(config)),
45 data_(nullptr) {
46 verifySettings_();
47
48 // Flat index doesn't need training
49 this->is_trained = true;
50
51 // Construct index
52 DeviceScope scope(device_);
53 data_ = new FlatIndex(resources,
54 dims,
55 metric == faiss::METRIC_L2,
56 config_.useFloat16,
57 config_.useFloat16Accumulator,
58 config_.storeTransposed,
59 memorySpace_);
60 }
61
62 GpuIndexFlat::~GpuIndexFlat() {
63 delete data_;
64 }
65
66 void
67 GpuIndexFlat::copyFrom(const faiss::IndexFlat* index) {
68 DeviceScope scope(device_);
69
70 this->d = index->d;
71 this->metric_type = index->metric_type;
72
73 // GPU code has 32 bit indices
74 FAISS_THROW_IF_NOT_FMT(index->ntotal <=
75 (faiss::Index::idx_t) std::numeric_limits<int>::max(),
76 "GPU index only supports up to %zu indices; "
77 "attempting to copy CPU index with %zu parameters",
78 (size_t) std::numeric_limits<int>::max(),
79 (size_t) index->ntotal);
80 this->ntotal = index->ntotal;
81
82 delete data_;
83 data_ = new FlatIndex(resources_,
84 this->d,
85 index->metric_type == faiss::METRIC_L2,
86 config_.useFloat16,
87 config_.useFloat16Accumulator,
88 config_.storeTransposed,
89 memorySpace_);
90
91 // The index could be empty
92 if (index->ntotal > 0) {
93 data_->add(index->xb.data(),
94 index->ntotal,
95 resources_->getDefaultStream(device_));
96 }
97 }
98
99 void
100 GpuIndexFlat::copyTo(faiss::IndexFlat* index) const {
101 DeviceScope scope(device_);
102
103 index->d = this->d;
104 index->ntotal = this->ntotal;
105 index->metric_type = this->metric_type;
106
107 FAISS_ASSERT(data_);
108 FAISS_ASSERT(data_->getSize() == this->ntotal);
109 index->xb.resize(this->ntotal * this->d);
110
111 auto stream = resources_->getDefaultStream(device_);
112
113 if (this->ntotal > 0) {
114 if (config_.useFloat16) {
115 auto vecFloat32 = data_->getVectorsFloat32Copy(stream);
116 fromDevice(vecFloat32, index->xb.data(), stream);
117 } else {
118 fromDevice(data_->getVectorsFloat32Ref(), index->xb.data(), stream);
119 }
120 }
121 }
122
123 size_t
124 GpuIndexFlat::getNumVecs() const {
125 return this->ntotal;
126 }
127
128 void
129 GpuIndexFlat::reset() {
130 DeviceScope scope(device_);
131
132 // Free the underlying memory
133 data_->reset();
134 this->ntotal = 0;
135 }
136
137 void
138 GpuIndexFlat::train(Index::idx_t n, const float* x) {
139 // nothing to do
140 }
141
142 void
143 GpuIndexFlat::add(Index::idx_t n, const float* x) {
144 FAISS_THROW_IF_NOT_MSG(this->is_trained, "Index not trained");
145
146 // For now, only support <= max int results
147 FAISS_THROW_IF_NOT_FMT(n <= (Index::idx_t) std::numeric_limits<int>::max(),
148 "GPU index only supports up to %d indices",
149 std::numeric_limits<int>::max());
150
151 if (n == 0) {
152 // nothing to add
153 return;
154 }
155
156 DeviceScope scope(device_);
157
158 // To avoid multiple re-allocations, ensure we have enough storage
159 // available
160 data_->reserve(n, resources_->getDefaultStream(device_));
161
162 // If we're not operating in float16 mode, we don't need the input
163 // data to be resident on our device; we can add directly.
164 if (!config_.useFloat16) {
165 addImpl_(n, x, nullptr);
166 } else {
167 // Otherwise, perform the paging
168 GpuIndex::add(n, x);
169 }
170 }
171
172 bool
173 GpuIndexFlat::addImplRequiresIDs_() const {
174 return false;
175 }
176
177 void
178 GpuIndexFlat::addImpl_(int n,
179 const float* x,
180 const Index::idx_t* ids) {
181 FAISS_ASSERT(data_);
182 FAISS_ASSERT(n > 0);
183
184 // We do not support add_with_ids
185 FAISS_THROW_IF_NOT_MSG(!ids, "add_with_ids not supported");
186
187 // Due to GPU indexing in int32, we can't store more than this
188 // number of vectors on a GPU
189 FAISS_THROW_IF_NOT_FMT(this->ntotal + n <=
190 (faiss::Index::idx_t) std::numeric_limits<int>::max(),
191 "GPU index only supports up to %zu indices",
192 (size_t) std::numeric_limits<int>::max());
193
194 data_->add(x, n, resources_->getDefaultStream(device_));
195 this->ntotal += n;
196 }
197
198 void
199 GpuIndexFlat::searchImpl_(int n,
200 const float* x,
201 int k,
202 float* distances,
203 Index::idx_t* labels) const {
204 auto stream = resources_->getDefaultStream(device_);
205
206 // Input and output data are already resident on the GPU
207 Tensor<float, 2, true> queries(const_cast<float*>(x), {n, (int) this->d});
208 Tensor<float, 2, true> outDistances(distances, {n, k});
209 Tensor<Index::idx_t, 2, true> outLabels(labels, {n, k});
210
211 // FlatIndex only supports int indices
212 DeviceTensor<int, 2, true> outIntLabels(
213 resources_->getMemoryManagerCurrentDevice(), {n, k}, stream);
214
215 data_->query(queries, k, outDistances, outIntLabels, true);
216
217 // Convert int to idx_t
218 thrust::transform(thrust::cuda::par.on(stream),
219 outIntLabels.data(),
220 outIntLabels.end(),
221 outLabels.data(),
222 IntToIdxType());
223 }
224
225 void
226 GpuIndexFlat::reconstruct(faiss::Index::idx_t key,
227 float* out) const {
228 DeviceScope scope(device_);
229
230 FAISS_THROW_IF_NOT_MSG(key < this->ntotal, "index out of bounds");
231 auto stream = resources_->getDefaultStream(device_);
232
233 if (config_.useFloat16) {
234 auto vec = data_->getVectorsFloat32Copy(key, 1, stream);
235 fromDevice(vec.data(), out, this->d, stream);
236 } else {
237 auto vec = data_->getVectorsFloat32Ref()[key];
238 fromDevice(vec.data(), out, this->d, stream);
239 }
240 }
241
242 void
243 GpuIndexFlat::reconstruct_n(faiss::Index::idx_t i0,
244 faiss::Index::idx_t num,
245 float* out) const {
246 DeviceScope scope(device_);
247
248 FAISS_THROW_IF_NOT_MSG(i0 < this->ntotal, "index out of bounds");
249 FAISS_THROW_IF_NOT_MSG(i0 + num - 1 < this->ntotal, "num out of bounds");
250 auto stream = resources_->getDefaultStream(device_);
251
252 if (config_.useFloat16) {
253 auto vec = data_->getVectorsFloat32Copy(i0, num, stream);
254 fromDevice(vec.data(), out, num * this->d, stream);
255 } else {
256 auto vec = data_->getVectorsFloat32Ref()[i0];
257 fromDevice(vec.data(), out, this->d * num, stream);
258 }
259 }
260
261 void
262 GpuIndexFlat::verifySettings_() const {
263 // If we want Hgemm, ensure that it is supported on this device
264 if (config_.useFloat16Accumulator) {
265 #ifdef FAISS_USE_FLOAT16
266 FAISS_THROW_IF_NOT_MSG(config_.useFloat16,
267 "useFloat16Accumulator can only be enabled "
268 "with useFloat16");
269
270 FAISS_THROW_IF_NOT_FMT(getDeviceSupportsFloat16Math(config_.device),
271 "Device %d does not support Hgemm "
272 "(useFloat16Accumulator)",
273 config_.device);
274 #else
275 FAISS_THROW_IF_NOT_MSG(false, "not compiled with float16 support");
276 #endif
277 }
278 }
279
280 //
281 // GpuIndexFlatL2
282 //
283
284 GpuIndexFlatL2::GpuIndexFlatL2(GpuResources* resources,
285 faiss::IndexFlatL2* index,
286 GpuIndexFlatConfig config) :
287 GpuIndexFlat(resources, index, config) {
288 }
289
290 GpuIndexFlatL2::GpuIndexFlatL2(GpuResources* resources,
291 int dims,
292 GpuIndexFlatConfig config) :
293 GpuIndexFlat(resources, dims, faiss::METRIC_L2, config) {
294 }
295
296 void
297 GpuIndexFlatL2::copyFrom(faiss::IndexFlatL2* index) {
298 GpuIndexFlat::copyFrom(index);
299 }
300
301 void
302 GpuIndexFlatL2::copyTo(faiss::IndexFlatL2* index) {
303 GpuIndexFlat::copyTo(index);
304 }
305
306 //
307 // GpuIndexFlatIP
308 //
309
310 GpuIndexFlatIP::GpuIndexFlatIP(GpuResources* resources,
311 faiss::IndexFlatIP* index,
312 GpuIndexFlatConfig config) :
313 GpuIndexFlat(resources, index, config) {
314 }
315
316 GpuIndexFlatIP::GpuIndexFlatIP(GpuResources* resources,
317 int dims,
318 GpuIndexFlatConfig config) :
319 GpuIndexFlat(resources, dims, faiss::METRIC_INNER_PRODUCT, config) {
320 }
321
322 void
323 GpuIndexFlatIP::copyFrom(faiss::IndexFlatIP* index) {
324 GpuIndexFlat::copyFrom(index);
325 }
326
327 void
328 GpuIndexFlatIP::copyTo(faiss::IndexFlatIP* index) {
329 GpuIndexFlat::copyTo(index);
330 }
331
332 } } // namespace
faiss::gpu::FlatIndex::getVectorsFloat32Copy
DeviceTensor< float, 2, true > getVectorsFloat32Copy(cudaStream_t stream)
Definition: FlatIndex.cu:89
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::GpuIndexFlatL2::copyFrom
void copyFrom(faiss::IndexFlatL2 *index)
Definition: GpuIndexFlat.cu:297
faiss::gpu::GpuIndexFlat::copyTo
void copyTo(faiss::IndexFlat *index) const
Definition: GpuIndexFlat.cu:100
faiss::gpu::GpuIndexFlat::reconstruct_n
void reconstruct_n(faiss::Index::idx_t i0, faiss::Index::idx_t num, float *out) const override
Batch reconstruction method.
Definition: GpuIndexFlat.cu:243
faiss::gpu::FlatIndex::getSize
int getSize() const
Returns the number of vectors we contain.
Definition: FlatIndex.cu:45
faiss::gpu::FlatIndex
Holder of GPU resources for a particular flat index.
Definition: FlatIndex.cuh:21
faiss::gpu::GpuIndexFlatL2::copyTo
void copyTo(faiss::IndexFlatL2 *index)
Definition: GpuIndexFlat.cu:302
faiss::IndexFlatL2
Definition: IndexFlat.h:77
faiss::gpu::IntToIdxType
Definition: ConversionOperators.cuh:21
faiss::gpu::GpuIndexFlat::getNumVecs
size_t getNumVecs() const
Returns the number of vectors we contain.
Definition: GpuIndexFlat.cu:124
faiss::gpu::GpuIndexFlat::searchImpl_
void searchImpl_(int n, const float *x, int k, float *distances, faiss::Index::idx_t *labels) const override
Called from GpuIndex for search.
Definition: GpuIndexFlat.cu:199
faiss::gpu::GpuResources
Definition: GpuResources.h:21
faiss::gpu::GpuIndexFlat::GpuIndexFlat
GpuIndexFlat(GpuResources *resources, const faiss::IndexFlat *index, GpuIndexFlatConfig config=GpuIndexFlatConfig())
Definition: GpuIndexFlat.cu:25
faiss::IndexFlatIP
Definition: IndexFlat.h:71
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::GpuIndex
Definition: GpuIndex.h:33
faiss::gpu::GpuIndexFlatL2::GpuIndexFlatL2
GpuIndexFlatL2(GpuResources *resources, faiss::IndexFlatL2 *index, GpuIndexFlatConfig config=GpuIndexFlatConfig())
Definition: GpuIndexFlat.cu:284
faiss::gpu::GpuResources::getDefaultStream
virtual cudaStream_t getDefaultStream(int device)=0
faiss::Index::d
int d
vector dimension
Definition: Index.h:66
faiss::Index::idx_t
long idx_t
all indices are this type
Definition: Index.h:62
faiss::gpu::GpuIndexFlatConfig
Definition: GpuIndexFlat.h:25
faiss::gpu::GpuIndexFlat::reconstruct
void reconstruct(faiss::Index::idx_t key, float *out) const override
Definition: GpuIndexFlat.cu:226
faiss::gpu::GpuResources::getMemoryManagerCurrentDevice
DeviceMemory & getMemoryManagerCurrentDevice()
Calls getMemoryManager for the current device.
Definition: GpuResources.cpp:33
faiss::gpu::GpuIndex::device_
const int device_
The GPU device we are resident on.
Definition: GpuIndex.h:126
faiss::gpu::GpuIndexFlatIP::copyTo
void copyTo(faiss::IndexFlatIP *index)
Definition: GpuIndexFlat.cu:328
faiss::gpu::GpuIndexFlatIP::GpuIndexFlatIP
GpuIndexFlatIP(GpuResources *resources, faiss::IndexFlatIP *index, GpuIndexFlatConfig config=GpuIndexFlatConfig())
Definition: GpuIndexFlat.cu:310
faiss::gpu::GpuIndex::resources_
GpuResources * resources_
Manages streams, cuBLAS handles and scratch memory for devices.
Definition: GpuIndex.h:123
faiss::gpu::GpuIndexFlat::addImpl_
void addImpl_(int n, const float *x, const Index::idx_t *ids) override
Called from GpuIndex for add.
Definition: GpuIndexFlat.cu:178
faiss::gpu::FlatIndex::reserve
void reserve(size_t numVecs, cudaStream_t stream)
Reserve storage that can contain at least this many vectors.
Definition: FlatIndex.cu:66
faiss::gpu::FlatIndex::add
void add(const float *data, int numVecs, cudaStream_t stream)
Definition: FlatIndex.cu:196
faiss::Index::ntotal
idx_t ntotal
total nb of indexed vectors
Definition: Index.h:67
faiss::gpu::GpuIndex::add
void add(faiss::Index::idx_t, const float *x) override
Definition: GpuIndex.cu:79
faiss::gpu::GpuIndexFlat::copyFrom
void copyFrom(const faiss::IndexFlat *index)
Definition: GpuIndexFlat.cu:67
faiss::gpu::Tensor
Our tensor type.
Definition: Tensor.cuh:28
faiss::gpu::DeviceScope
Definition: DeviceUtils.h:68
faiss::gpu::GpuIndex::memorySpace_
const MemorySpace memorySpace_
The memory space of our primary storage on the GPU.
Definition: GpuIndex.h:129
faiss::gpu::GpuIndexFlat::addImplRequiresIDs_
bool addImplRequiresIDs_() const override
Flat index does not require IDs as there is no storage available for them.
Definition: GpuIndexFlat.cu:173
faiss::gpu::GpuIndexFlat::config_
const GpuIndexFlatConfig config_
Our config object.
Definition: GpuIndexFlat.h:123
faiss::Index::metric_type
MetricType metric_type
type of metric this index uses for search
Definition: Index.h:74
faiss::gpu::GpuIndexFlat::add
void add(faiss::Index::idx_t, const float *x) override
Overrides to avoid excessive copies.
Definition: GpuIndexFlat.cu:143
faiss::gpu::FlatIndex::getVectorsFloat32Ref
Tensor< float, 2, true > & getVectorsFloat32Ref()
Returns a reference to our vectors currently in use.
Definition: FlatIndex.cu:77
faiss::gpu::GpuIndexFlat
Definition: GpuIndexFlat.h:52
faiss::gpu::GpuIndexFlat::reset
void reset() override
Clears all vectors from this index.
Definition: GpuIndexFlat.cu:129
faiss::gpu::GpuIndexFlatIP::copyFrom
void copyFrom(faiss::IndexFlatIP *index)
Definition: GpuIndexFlat.cu:323
faiss::Index::is_trained
bool is_trained
set if the Index does not require training, or if training is done already
Definition: Index.h:71
faiss::gpu::FlatIndex::reset
void reset()
Free all storage.
Definition: FlatIndex.cu:270
faiss::gpu::DeviceTensor
Definition: DeviceTensor.cuh:22
faiss::IndexFlat::xb
std::vector< float > xb
database vectors, size ntotal * d
Definition: IndexFlat.h:23
faiss::gpu::GpuIndexFlat::train
void train(Index::idx_t n, const float *x) override
This index is not trained, so this does nothing.
Definition: GpuIndexFlat.cu:138
faiss::MetricType
MetricType
Some algorithms support both an inner product version and a L2 search version.
Definition: Index.h:44
faiss::gpu::GpuIndexFlat::data_
FlatIndex * data_
Definition: GpuIndexFlat.h:127