Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

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