Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
GpuIndexIVFFlat.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 "GpuIndexIVFFlat.h"
13 #include "../IndexFlat.h"
14 #include "../IndexIVF.h"
15 #include "GpuIndexFlat.h"
16 #include "GpuResources.h"
17 #include "impl/IVFFlat.cuh"
18 #include "utils/CopyUtils.cuh"
19 #include "utils/DeviceUtils.h"
20 #include "utils/Float16.cuh"
21
22 #include <limits>
23
24 namespace faiss { namespace gpu {
25
26 GpuIndexIVFFlat::GpuIndexIVFFlat(GpuResources* resources,
27 int device,
28 bool useFloat16CoarseQuantizer,
29 bool useFloat16IVFStorage,
30 int dims,
31 int nlist,
32 IndicesOptions indicesOptions,
33 faiss::MetricType metric) :
34 GpuIndexIVF(resources,
35 device,
36 indicesOptions,
37 useFloat16CoarseQuantizer,
38 dims,
39 metric,
40 nlist),
41 useFloat16IVFStorage_(useFloat16IVFStorage),
42 reserveMemoryVecs_(0),
43 index_(nullptr) {
44 DeviceScope scope(device_);
45
46 // faiss::Index params
47 this->is_trained = false;
48
49 #ifndef FAISS_USE_FLOAT16
50 FAISS_ASSERT(!useFloat16IVFStorage_,
51 "float16 unsupported; need CUDA SDK >= 7.5");
52 #endif
53
54 // We haven't trained ourselves, so don't construct the IVFFlat
55 // index yet
56 }
57
58 GpuIndexIVFFlat::GpuIndexIVFFlat(GpuResources* resources,
59 int device,
60 GpuIndexFlat* quantizer,
61 bool useFloat16IVFStorage,
62 int dims,
63 int nlist,
64 IndicesOptions indicesOptions,
65 faiss::MetricType metric) :
66 GpuIndexIVF(resources, device, indicesOptions, dims,
67 metric, nlist, quantizer),
68 useFloat16IVFStorage_(useFloat16IVFStorage),
69 reserveMemoryVecs_(0),
70 index_(nullptr) {
71 DeviceScope scope(device_);
72
73 // faiss::Index params
74 this->is_trained = (quantizer->ntotal > 0);
75
76 #ifndef FAISS_USE_FLOAT16
77 FAISS_ASSERT(!useFloat16IVFStorage_,
78 "float16 unsupported; need CUDA SDK >= 7.5");
79 #endif
80
81 if (this->is_trained) {
82 index_ = new IVFFlat(resources_,
83 quantizer_->getGpuData(),
84 quantizer_->metric_type == faiss::METRIC_L2,
85 useFloat16IVFStorage_,
86 indicesOptions_);
87 }
88 }
89
90 GpuIndexIVFFlat::~GpuIndexIVFFlat() {
91 delete index_;
92 }
93
94 void
95 GpuIndexIVFFlat::reserveMemory(size_t numVecs) {
96 reserveMemoryVecs_ = numVecs;
97 if (index_) {
98 index_->reserveMemory(numVecs);
99 }
100 }
101
102 void
103 GpuIndexIVFFlat::copyFrom(const faiss::IndexIVFFlat* index) {
104 DeviceScope scope(device_);
105
106 GpuIndexIVF::copyFrom(index);
107
108 // Clear out our old data
109 delete index_;
110 index_ = nullptr;
111
112 // The other index might not be trained
113 if (!index->is_trained) {
114 return;
115 }
116
117 // Otherwise, we can populate ourselves from the other index
118 this->is_trained = true;
119
120 // Copy our lists as well
121 index_ = new IVFFlat(resources_,
122 quantizer_->getGpuData(),
123 index->metric_type == faiss::METRIC_L2,
124 useFloat16IVFStorage_,
125 indicesOptions_);
126
127 FAISS_ASSERT(index->vecs.size() == index->ids.size());
128 for (size_t i = 0; i < index->vecs.size(); ++i) {
129 auto& vecs = index->vecs[i];
130 auto& ids = index->ids[i];
131
132 FAISS_ASSERT(vecs.size() % this->d == 0);
133 auto numVecs = vecs.size() / this->d;
134 FAISS_ASSERT(numVecs == ids.size());
135
136 index_->addCodeVectorsFromCpu(i, vecs.data(), ids.data(), numVecs);
137 }
138 }
139
140 void
141 GpuIndexIVFFlat::copyTo(faiss::IndexIVFFlat* index) const {
142 DeviceScope scope(device_);
143
144 // We must have the indices in order to copy to ourselves
145 FAISS_ASSERT(indicesOptions_ != INDICES_IVF);
146
147 GpuIndexIVF::copyTo(index);
148
149 // Clear out the old inverted lists
150 index->vecs.clear();
151 index->vecs.resize(nlist_);
152
153 // Copy the inverted lists
154 if (index_) {
155 for (int i = 0; i < nlist_; ++i) {
156 index->vecs[i] = index_->getListVectors(i);
157 index->ids[i] = index_->getListIndices(i);
158 }
159 }
160 }
161
162 size_t
163 GpuIndexIVFFlat::reclaimMemory() {
164 if (index_) {
165 DeviceScope scope(device_);
166
167 return index_->reclaimMemory();
168 }
169
170 return 0;
171 }
172
173 void
174 GpuIndexIVFFlat::reset() {
175 if (index_) {
176 DeviceScope scope(device_);
177
178 index_->reset();
179 this->ntotal = 0;
180 } else {
181 FAISS_ASSERT(this->ntotal == 0);
182 }
183 }
184
185 void
186 GpuIndexIVFFlat::train(Index::idx_t n, const float* x) {
187 DeviceScope scope(device_);
188
189 if (this->is_trained) {
190 FAISS_ASSERT(quantizer_->is_trained);
191 FAISS_ASSERT(quantizer_->ntotal == nlist_);
192 FAISS_ASSERT(index_);
193 return;
194 }
195
196 FAISS_ASSERT(!index_);
197
198 trainQuantizer_(n, x);
199
200 // The quantizer is now trained; construct the IVF index
201 index_ = new IVFFlat(resources_,
202 quantizer_->getGpuData(),
203 this->metric_type == faiss::METRIC_L2,
204 useFloat16IVFStorage_,
205 indicesOptions_);
206 if (reserveMemoryVecs_) {
207 index_->reserveMemory(reserveMemoryVecs_);
208 }
209
210 this->is_trained = true;
211 }
212
213 void
214 GpuIndexIVFFlat::addImpl_(Index::idx_t n,
215 const float* x,
216 const Index::idx_t* xids) {
217 // Device is already set in GpuIndex::addInternal_
218
219 FAISS_ASSERT(index_);
220 FAISS_ASSERT(n > 0);
221
222 auto stream = resources_->getDefaultStreamCurrentDevice();
223
224 auto deviceVecs =
225 toDevice<float, 2>(resources_,
226 device_,
227 const_cast<float*>(x),
228 stream,
229 {(int) n, index_->getDim()});
230
231 static_assert(sizeof(long) == sizeof(Index::idx_t), "size mismatch");
232 auto deviceIds =
233 toDevice<long, 1>(resources_,
234 device_,
235 const_cast<long*>(xids),
236 stream,
237 {(int) n});
238
239 // Not all vectors may be able to be added (some may contain NaNs
240 // etc)
241 ntotal += index_->classifyAndAddVectors(deviceVecs, deviceIds);
242 }
243
244 void
245 GpuIndexIVFFlat::searchImpl_(faiss::Index::idx_t n,
246 const float* x,
247 faiss::Index::idx_t k,
248 float* distances,
249 faiss::Index::idx_t* labels) const {
250 // Device is already set in GpuIndex::search
251
252 FAISS_ASSERT(index_);
253 FAISS_ASSERT(n > 0);
254
255 auto stream = resources_->getDefaultStream(device_);
256
257 // Make sure arguments are on the device we desire; use temporary
258 // memory allocations to move it if necessary
259 auto devX =
260 toDevice<float, 2>(resources_,
261 device_,
262 const_cast<float*>(x),
263 stream,
264 {(int) n, this->d});
265 auto devDistances =
266 toDevice<float, 2>(resources_,
267 device_,
268 distances,
269 stream,
270 {(int) n, (int) k});
271 auto devLabels =
272 toDevice<faiss::Index::idx_t, 2>(resources_,
273 device_,
274 labels,
275 stream,
276 {(int) n, (int) k});
277
278 index_->query(devX, nprobe_, k, devDistances, devLabels);
279
280 // Copy back if necessary
281 fromDevice<float, 2>(devDistances, distances, stream);
282 fromDevice<faiss::Index::idx_t, 2>(devLabels, labels, stream);
283 }
284
285 void
286 GpuIndexIVFFlat::set_typename() {
287 this->index_typename = "GpuIndexIVFFlat";
288 }
289
290 } } // namespace
faiss::gpu::IVFFlat
Definition: IVFFlat.cuh:18
faiss::gpu::IVFBase::getDim
int getDim() const
Return the number of dimensions we are indexing.
Definition: IVFBase.cu:99
faiss::gpu::GpuIndexFlat::getGpuData
FlatIndex * getGpuData()
For internal access.
Definition: GpuIndexFlat.h:109
faiss::gpu::GpuIndexIVFFlat::copyFrom
void copyFrom(const faiss::IndexIVFFlat *index)
Definition: GpuIndexIVFFlat.cu:103
faiss::gpu::IVFBase::reserveMemory
void reserveMemory(size_t numVecs)
Reserve GPU memory in our inverted lists for this number of vectors.
Definition: IVFBase.cu:44
faiss::gpu::GpuIndexIVFFlat::reclaimMemory
size_t reclaimMemory()
Definition: GpuIndexIVFFlat.cu:163
faiss::gpu::GpuResources
Definition: GpuResources.h:24
faiss::IndexIVFFlat
Definition: IndexIVF.h:128
faiss::gpu::GpuIndex::device_
int device_
The GPU device we are resident on.
Definition: GpuIndex.h:80
faiss::IndexIVF::ids
std::vector< std::vector< long > > ids
Inverted lists for indexes.
Definition: IndexIVF.h:56
faiss::gpu::GpuIndexIVF
Definition: GpuIndexIVF.h:25
faiss::Index::d
int d
vector dimension
Definition: Index.h:66
faiss::gpu::IVFFlat::classifyAndAddVectors
int classifyAndAddVectors(Tensor< float, 2, true > &vecs, Tensor< long, 1, true > &indices)
Definition: IVFFlat.cu:129
faiss::gpu::GpuIndexIVFFlat::copyTo
void copyTo(faiss::IndexIVFFlat *index) const
Definition: GpuIndexIVFFlat.cu:141
faiss::gpu::GpuIndexIVF::nprobe_
int nprobe_
Number of inverted list probes per query.
Definition: GpuIndexIVF.h:92
faiss::gpu::GpuIndexIVF::indicesOptions_
const IndicesOptions indicesOptions_
How should indices be stored on the GPU?
Definition: GpuIndexIVF.h:82
faiss::gpu::GpuIndexIVFFlat::reserveMemory
void reserveMemory(size_t numVecs)
Reserve GPU memory in our inverted lists for this number of vectors.
Definition: GpuIndexIVFFlat.cu:95
faiss::gpu::GpuIndex::resources_
GpuResources * resources_
Manages streans, cuBLAS handles and scratch memory for devices.
Definition: GpuIndex.h:77
faiss::gpu::GpuIndexIVF::copyTo
void copyTo(faiss::IndexIVF *index) const
Copy what we have to the CPU equivalent.
Definition: GpuIndexIVF.cu:187
faiss::Index::idx_t
long idx_t
all indices are this type
Definition: Index.h:64
faiss::gpu::GpuIndexIVF::nlist_
int nlist_
Number of inverted lists that we manage.
Definition: GpuIndexIVF.h:89
faiss::gpu::GpuIndexIVFFlat::searchImpl_
virtual void searchImpl_(faiss::Index::idx_t n, const float *x, faiss::Index::idx_t k, float *distances, faiss::Index::idx_t *labels) const
Called from GpuIndex for search.
Definition: GpuIndexIVFFlat.cu:245
faiss::Index::ntotal
idx_t ntotal
total nb of indexed vectors
Definition: Index.h:67
faiss::gpu::GpuIndexIVFFlat::reset
virtual void reset()
removes all elements from the database.
Definition: GpuIndexIVFFlat.cu:174
faiss::gpu::IVFBase::reset
void reset()
Definition: IVFBase.cu:74
faiss::gpu::GpuIndexIVFFlat::train
virtual void train(Index::idx_t n, const float *x)
Definition: GpuIndexIVFFlat.cu:186
faiss::gpu::IVFFlat::addCodeVectorsFromCpu
void addCodeVectorsFromCpu(int listId, const float *vecs, const long *indices, size_t numVecs)
Definition: IVFFlat.cu:57
faiss::gpu::DeviceScope
Definition: DeviceUtils.h:55
faiss::gpu::GpuIndexIVF::quantizer_
GpuIndexFlat * quantizer_
Quantizer for inverted lists.
Definition: GpuIndexIVF.h:98
faiss::gpu::IVFFlat::query
void query(Tensor< float, 2, true > &queries, int nprobe, int k, Tensor< float, 2, true > &outDistances, Tensor< long, 2, true > &outIndices)
Definition: IVFFlat.cu:287
faiss::gpu::IVFFlat::getListVectors
std::vector< float > getListVectors(int listId) const
Return the vectors of a particular list back to the CPU.
Definition: IVFFlat.cu:353
faiss::Index::metric_type
MetricType metric_type
type of metric this index uses for search
Definition: Index.h:74
faiss::gpu::GpuIndexIVFFlat::addImpl_
virtual void addImpl_(faiss::Index::idx_t n, const float *x, const faiss::Index::idx_t *ids)
Called from GpuIndex for add/add_with_ids.
Definition: GpuIndexIVFFlat.cu:214
faiss::gpu::GpuIndexIVFFlat::GpuIndexIVFFlat
GpuIndexIVFFlat(GpuResources *resources, int device, bool useFloat16CoarseQuantizer, bool useFloat16IVFStorage, int dims, int nlist, IndicesOptions indicesOptions, faiss::MetricType metric)
Definition: GpuIndexIVFFlat.cu:26
faiss::gpu::IVFBase::getListIndices
std::vector< long > getListIndices(int listId) const
Return the list indices of a particular list back to the CPU.
Definition: IVFBase.cu:206
faiss::gpu::GpuIndexFlat
Definition: GpuIndexFlat.h:43
faiss::gpu::GpuIndexIVF::copyFrom
void copyFrom(const faiss::IndexIVF *index)
Copy what we need from the CPU equivalent.
Definition: GpuIndexIVF.cu:120
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::IVFBase::reclaimMemory
size_t reclaimMemory()
Definition: IVFBase.cu:104
faiss::MetricType
MetricType
Some algorithms support both an inner product vetsion and a L2 search version.
Definition: Index.h:44
faiss::IndexIVFFlat::vecs
std::vector< std::vector< float > > vecs
Definition: IndexIVF.h:131