Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

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