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::add_with_ids(Index::idx_t n,
215 const float* x,
216 const Index::idx_t* xids) {
217 FAISS_ASSERT(this->is_trained);
218 FAISS_ASSERT(index_);
219
220 DeviceScope scope(device_);
221 auto stream = resources_->getDefaultStreamCurrentDevice();
222
223 auto deviceVecs =
224 toDevice<float, 2>(resources_,
225 device_,
226 const_cast<float*>(x),
227 stream,
228 {(int) n, index_->getDim()});
229
230 static_assert(sizeof(long) == sizeof(Index::idx_t), "size mismatch");
231 auto deviceIds =
232 toDevice<long, 1>(resources_,
233 device_,
234 const_cast<long*>(xids),
235 stream,
236 {(int) n});
237
238 // Not all vectors may be able to be added (some may contain NaNs
239 // etc)
240 ntotal += index_->classifyAndAddVectors(deviceVecs, deviceIds);
241 }
242
243 void
244 GpuIndexIVFFlat::search(faiss::Index::idx_t n,
245 const float* x,
246 faiss::Index::idx_t k,
247 float* distances,
248 faiss::Index::idx_t* labels) const {
249 if (n == 0) {
250 return;
251 }
252
253 FAISS_ASSERT(this->is_trained);
254 FAISS_ASSERT(index_);
255
256 DeviceScope scope(device_);
257 auto stream = resources_->getDefaultStream(device_);
258
259 // Make sure arguments are on the device we desire; use temporary
260 // memory allocations to move it if necessary
261 auto devX =
262 toDevice<float, 2>(resources_,
263 device_,
264 const_cast<float*>(x),
265 stream,
266 {(int) n, this->d});
267 auto devDistances =
268 toDevice<float, 2>(resources_,
269 device_,
270 distances,
271 stream,
272 {(int) n, (int) k});
273 auto devLabels =
274 toDevice<faiss::Index::idx_t, 2>(resources_,
275 device_,
276 labels,
277 stream,
278 {(int) n, (int) k});
279
280 index_->query(devX, nprobe_, k, devDistances, devLabels);
281
282 // Copy back if necessary
283 fromDevice<float, 2>(devDistances, distances, stream);
284 fromDevice<faiss::Index::idx_t, 2>(devLabels, labels, stream);
285 }
286
287 void
288 GpuIndexIVFFlat::set_typename() {
289 this->index_typename = "GpuIndexIVFFlat";
290 }
291
292 } } // 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:100
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:113
faiss::gpu::GpuIndexIVFFlat::train
void train(Index::idx_t n, const float *x) override
Definition: GpuIndexIVFFlat.cu:186
faiss::gpu::GpuIndex::device_
int device_
The GPU device we are resident on.
Definition: GpuIndex.h:43
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:91
faiss::gpu::GpuIndexIVF::indicesOptions_
const IndicesOptions indicesOptions_
How should indices be stored on the GPU?
Definition: GpuIndexIVF.h:81
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:40
faiss::gpu::GpuIndexIVF::copyTo
void copyTo(faiss::IndexIVF *index) const
Copy what we have to the CPU equivalent.
Definition: GpuIndexIVF.cu:181
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:88
faiss::Index::ntotal
idx_t ntotal
total nb of indexed vectors
Definition: Index.h:67
faiss::gpu::IVFBase::reset
void reset()
Definition: IVFBase.cu:74
faiss::gpu::GpuIndexIVFFlat::add_with_ids
void add_with_ids(Index::idx_t n, const float *x, const Index::idx_t *xids) override
Definition: GpuIndexIVFFlat.cu:214
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:97
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::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:31
faiss::gpu::GpuIndexIVF::copyFrom
void copyFrom(const faiss::IndexIVF *index)
Copy what we need from the CPU equivalent.
Definition: GpuIndexIVF.cu:117
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::gpu::GpuIndexIVFFlat::reset
void reset() override
removes all elements from the database.
Definition: GpuIndexIVFFlat.cu:174
faiss::gpu::GpuIndexIVFFlat::search
void search(faiss::Index::idx_t n, const float *x, faiss::Index::idx_t k, float *distances, faiss::Index::idx_t *labels) const override
Definition: GpuIndexIVFFlat.cu:244
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:116