Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

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