Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
GpuIndexIVF.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 "GpuIndexIVF.h"
13 #include "../FaissAssert.h"
14 #include "../IndexFlat.h"
15 #include "../IndexIVF.h"
16 #include "GpuIndexFlat.h"
17 #include "utils/DeviceUtils.h"
18 #include "utils/Float16.cuh"
19
20 namespace faiss { namespace gpu {
21
22 GpuIndexIVF::GpuIndexIVF(GpuResources* resources,
23 int device,
24 IndicesOptions indicesOptions,
25 bool useFloat16CoarseQuantizer,
26 int dims,
27 faiss::MetricType metric,
28 int nlist) :
29 GpuIndex(resources, device, dims, metric),
30 indicesOptions_(indicesOptions),
31 useFloat16CoarseQuantizer_(useFloat16CoarseQuantizer),
32 nlist_(nlist),
33 nprobe_(1),
34 quantizer_(nullptr),
35 ownsQuantizer_(true) {
36 #ifndef FAISS_USE_FLOAT16
37 FAISS_ASSERT(!useFloat16CoarseQuantizer_,
38 "float16 unsupported; need CUDA SDK >= 7.5");
39 #endif
40
41 init_();
42 }
43
44 GpuIndexIVF::GpuIndexIVF(GpuResources* resources,
45 int device,
46 IndicesOptions indicesOptions,
47 int dims,
48 faiss::MetricType metric,
49 int nlist,
50 GpuIndexFlat* quantizer) :
51 GpuIndex(resources, device, dims, metric),
52 indicesOptions_(indicesOptions),
53 useFloat16CoarseQuantizer_(quantizer->getUseFloat16()),
54 nlist_(nlist),
55 nprobe_(1),
56 quantizer_(quantizer),
57 ownsQuantizer_(false) {
58 #ifndef FAISS_USE_FLOAT16
59 FAISS_ASSERT(!useFloat16CoarseQuantizer_,
60 "float16 unsupported; need CUDA SDK >= 7.5");
61 #endif
62
63 FAISS_ASSERT(quantizer_->d == this->d);
64 FAISS_ASSERT(quantizer_->metric_type == this->metric_type);
65
66 init_();
67 }
68
69 void
70 GpuIndexIVF::init_() {
71 FAISS_ASSERT(nlist_ > 0);
72
73 // Spherical by default if the metric is inner_product
74 if (this->metric_type == faiss::METRIC_INNER_PRODUCT) {
75 cp_.spherical = true;
76 }
77
78 // here we set a low # iterations because this is typically used
79 // for large clusterings
80 cp_.niter = 10;
81 cp_.verbose = this->verbose;
82
83 if (!quantizer_) {
84 // Construct an empty quantizer
85 GpuIndexFlatConfig config;
86 config.device = device_;
87 config.useFloat16 = useFloat16CoarseQuantizer_;
88 config.storeTransposed = false;
89
90 if (this->metric_type == faiss::METRIC_L2) {
91 // FIXME: 2 different float16 options?
92 quantizer_ = new GpuIndexFlatL2(resources_, this->d, config);
93 } else if (this->metric_type == faiss::METRIC_INNER_PRODUCT) {
94 // FIXME: 2 different float16 options?
95 quantizer_ = new GpuIndexFlatIP(resources_, this->d, config);
96 } else {
97 // unknown metric type
98 FAISS_ASSERT(false);
99 }
100 }
101 }
102
103 GpuIndexIVF::~GpuIndexIVF() {
104 if (ownsQuantizer_) {
105 delete quantizer_;
106 }
107 }
108
109 IndicesOptions
110 GpuIndexIVF::getIndicesOptions() const {
111 return indicesOptions_;
112 }
113
114 bool
115 GpuIndexIVF::getUseFloat16CoarseQuantizer() const {
116 return useFloat16CoarseQuantizer_;
117 }
118
119 void
120 GpuIndexIVF::copyFrom(const faiss::IndexIVF* index) {
121 DeviceScope scope(device_);
122
123 this->d = index->d;
124 this->metric_type = index->metric_type;
125
126 FAISS_ASSERT(index->nlist > 0);
127 FAISS_ASSERT(index->nlist <=
128 (faiss::Index::idx_t) std::numeric_limits<int>::max());
129 nlist_ = index->nlist;
130 nprobe_ = index->nprobe;
131
132 // The metric type may have changed as well, so we might have to
133 // change our quantizer
134 delete quantizer_;
135 quantizer_ = nullptr;
136
137 GpuIndexFlatConfig config;
138 config.device = device_;
139 config.useFloat16 = useFloat16CoarseQuantizer_;
140 config.storeTransposed = false;
141
142 if (index->metric_type == faiss::METRIC_L2) {
143 // FIXME: 2 different float16 options?
144 quantizer_ = new GpuIndexFlatL2(resources_, this->d, config);
145 } else if (index->metric_type == faiss::METRIC_INNER_PRODUCT) {
146 // FIXME: 2 different float16 options?
147 quantizer_ = new GpuIndexFlatIP(resources_, this->d, config);
148 } else {
149 // unknown metric type
150 FAISS_ASSERT(false);
151 }
152
153 if (!index->is_trained) {
154 this->is_trained = false;
155 this->ntotal = 0;
156 return;
157 }
158
159 // Otherwise, we can populate ourselves from the other index
160 this->is_trained = true;
161
162 // Only use `int` on GPU
163 FAISS_ASSERT(index->ntotal <=
164 (faiss::Index::idx_t) std::numeric_limits<int>::max());
165 this->ntotal = index->ntotal;
166
167 // Since we're trained, the quantizer must have data
168 FAISS_ASSERT(index->quantizer->ntotal > 0);
169
170 if (index->metric_type == faiss::METRIC_L2) {
171 auto q = dynamic_cast<faiss::IndexFlatL2*>(index->quantizer);
172 FAISS_ASSERT(q);
173
174 quantizer_->copyFrom(q);
175 } else if (index->metric_type == faiss::METRIC_INNER_PRODUCT) {
176 auto q = dynamic_cast<faiss::IndexFlatIP*>(index->quantizer);
177 FAISS_ASSERT(q);
178
179 quantizer_->copyFrom(q);
180 } else {
181 // unknown metric type
182 FAISS_ASSERT(false);
183 }
184 }
185
186 void
187 GpuIndexIVF::copyTo(faiss::IndexIVF* index) const {
188 DeviceScope scope(device_);
189
190 //
191 // Index information
192 //
193 index->ntotal = this->ntotal;
194 index->d = this->d;
195 index->metric_type = this->metric_type;
196 index->is_trained = this->is_trained;
197
198 //
199 // IndexIVF information
200 //
201 index->nlist = nlist_;
202 index->nprobe = nprobe_;
203
204 // Construct and copy the appropriate quantizer
205 faiss::IndexFlat* q = nullptr;
206
207 if (this->metric_type == faiss::METRIC_L2) {
208 q = new faiss::IndexFlatL2(this->d);
209
210 } else if (this->metric_type == faiss::METRIC_INNER_PRODUCT) {
211 q = new faiss::IndexFlatIP(this->d);
212
213 } else {
214 // unknown metric type
215 FAISS_ASSERT(false);
216 }
217
218 FAISS_ASSERT(quantizer_);
219 quantizer_->copyTo(q);
220
221 if (index->own_fields) {
222 delete index->quantizer;
223 }
224
225 index->quantizer = q;
226 index->quantizer_trains_alone = false;
227 index->own_fields = true;
228 index->cp = cp_;
229 index->ids.clear();
230 index->ids.resize(nlist_);
231 index->maintain_direct_map = false;
232 index->direct_map.clear();
233 }
234
235 int
236 GpuIndexIVF::getNumLists() const {
237 return nlist_;
238 }
239
240 void
241 GpuIndexIVF::setNumProbes(int nprobe) {
242 FAISS_ASSERT(nprobe > 0);
243 nprobe_ = nprobe;
244 }
245
246 int
247 GpuIndexIVF::getNumProbes() const {
248 return nprobe_;
249 }
250
251 void
252 GpuIndexIVF::add(Index::idx_t n, const float* x) {
253 // FIXME: GPU-ize
254 std::vector<Index::idx_t> ids(n);
255 for (Index::idx_t i = 0; i < n; ++i) {
256 ids[i] = this->ntotal + i;
257 }
258
259 add_with_ids(n, x, ids.data());
260 }
261
262 void
263 GpuIndexIVF::trainQuantizer_(faiss::Index::idx_t n, const float* x) {
264 if (n == 0) {
265 // nothing to do
266 return;
267 }
268
269 if (quantizer_->is_trained && (quantizer_->ntotal == nlist_)) {
270 if (this->verbose) {
271 printf ("IVF quantizer does not need training.\n");
272 }
273
274 return;
275 }
276
277 if (this->verbose) {
278 printf ("Training IVF quantizer on %ld vectors in %dD\n", n, d);
279 }
280
281 DeviceScope scope(device_);
282
283 // leverage the CPU-side k-means code, which works for the GPU
284 // flat index as well
285 quantizer_->reset();
286 Clustering clus(this->d, nlist_, cp_);
287 clus.verbose = verbose;
288 clus.train(n, x, *quantizer_);
289 quantizer_->is_trained = true;
290
291 FAISS_ASSERT(quantizer_->ntotal == nlist_);
292 }
293
294
295 } } // namespace
faiss::IndexIVF
Definition: IndexIVF.h:46
faiss::gpu::GpuIndexIVF::getIndicesOptions
IndicesOptions getIndicesOptions() const
What indices storage options are we using?
Definition: GpuIndexIVF.cu:110
faiss::gpu::GpuIndexIVF::getNumProbes
int getNumProbes() const
Returns our current number of list probes per query.
Definition: GpuIndexIVF.cu:247
faiss::IndexFlat
Definition: IndexFlat.h:24
faiss::gpu::GpuIndexIVF::setNumProbes
void setNumProbes(int nprobe)
Sets the number of list probes per query.
Definition: GpuIndexIVF.cu:241
faiss::ClusteringParameters::niter
int niter
clustering iterations
Definition: Clustering.h:26
faiss::gpu::GpuIndexIVF::useFloat16CoarseQuantizer_
const bool useFloat16CoarseQuantizer_
Definition: GpuIndexIVF.h:86
faiss::gpu::GpuIndexIVF::getNumLists
int getNumLists() const
Returns the number of inverted lists we&#39;re managing.
Definition: GpuIndexIVF.cu:236
faiss::gpu::GpuIndexFlat::copyTo
void copyTo(faiss::IndexFlat *index) const
Definition: GpuIndexFlat.cu:111
faiss::IndexIVF::nprobe
size_t nprobe
number of probes at query time
Definition: IndexIVF.h:48
faiss::gpu::GpuIndexIVF::getUseFloat16CoarseQuantizer
bool getUseFloat16CoarseQuantizer() const
Is our coarse quantizer storing and performing math in float16?
Definition: GpuIndexIVF.cu:115
faiss::IndexIVF::quantizer_trains_alone
bool quantizer_trains_alone
just pass over the trainset to quantizer
Definition: IndexIVF.h:51
faiss::IndexFlatL2
Definition: IndexFlat.h:81
faiss::IndexFlatIP
Definition: IndexFlat.h:75
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::GpuIndexFlat::reset
virtual void reset()
Clears all vectors from this index.
Definition: GpuIndexFlat.cu:139
faiss::Index::d
int d
vector dimension
Definition: Index.h:66
faiss::IndexIVF::quantizer
Index * quantizer
quantizer that maps vectors to inverted lists
Definition: IndexIVF.h:50
faiss::gpu::GpuIndexFlatConfig
Definition: GpuIndexFlat.h:28
faiss::IndexIVF::cp
ClusteringParameters cp
to override default clustering params
Definition: IndexIVF.h:54
faiss::gpu::GpuIndexIVF::nprobe_
int nprobe_
Number of inverted list probes per query.
Definition: GpuIndexIVF.h:92
faiss::IndexIVF::own_fields
bool own_fields
whether object owns the quantizer
Definition: IndexIVF.h:52
faiss::gpu::GpuIndexIVF::indicesOptions_
const IndicesOptions indicesOptions_
How should indices be stored on the GPU?
Definition: GpuIndexIVF.h:82
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::Index::ntotal
idx_t ntotal
total nb of indexed vectors
Definition: Index.h:67
faiss::gpu::GpuIndex::add_with_ids
virtual void add_with_ids(Index::idx_t n, const float *x, const Index::idx_t *ids)
Definition: GpuIndex.cu:42
faiss::Index::verbose
bool verbose
verbosity level
Definition: Index.h:68
faiss::gpu::GpuIndexFlatIP
Definition: GpuIndexFlat.h:176
faiss::gpu::GpuIndexFlat::copyFrom
void copyFrom(const faiss::IndexFlat *index)
Definition: GpuIndexFlat.cu:84
faiss::gpu::DeviceScope
Definition: DeviceUtils.h:55
faiss::gpu::GpuIndexIVF::quantizer_
GpuIndexFlat * quantizer_
Quantizer for inverted lists.
Definition: GpuIndexIVF.h:98
faiss::gpu::GpuIndexFlatL2
Definition: GpuIndexFlat.h:151
faiss::Index::metric_type
MetricType metric_type
type of metric this index uses for search
Definition: Index.h:74
faiss::IndexIVF::nlist
size_t nlist
number of possible key values
Definition: IndexIVF.h:47
faiss::gpu::GpuIndexIVF::cp_
ClusteringParameters cp_
Ability to override default clustering parameters.
Definition: GpuIndexIVF.h:95
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::GpuIndexIVF::ownsQuantizer_
bool ownsQuantizer_
Do we own the above quantizer instance?
Definition: GpuIndexIVF.h:101
faiss::gpu::GpuIndexIVF::add
void add(Index::idx_t n, const float *x) override
Definition: GpuIndexIVF.cu:252
faiss::IndexIVF::maintain_direct_map
bool maintain_direct_map
map for direct access to the elements. Enables reconstruct().
Definition: IndexIVF.h:59
faiss::ClusteringParameters::spherical
bool spherical
do we want normalized centroids?
Definition: Clustering.h:30
faiss::MetricType
MetricType
Some algorithms support both an inner product vetsion and a L2 search version.
Definition: Index.h:44