Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

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