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 #include <faiss/IndexFlat.h>
9 #include <faiss/IndexIVF.h>
10 #include <faiss/gpu/GpuIndexFlat.h>
11 #include <faiss/gpu/GpuIndexIVF.h>
12 #include <faiss/gpu/utils/DeviceUtils.h>
13 #include <faiss/impl/FaissAssert.h>
14 #include <faiss/gpu/utils/Float16.cuh>
15
16 namespace faiss {
17 namespace gpu {
18
GpuIndexIVF(GpuResourcesProvider * provider,int dims,faiss::MetricType metric,float metricArg,int nlistIn,GpuIndexIVFConfig config)19 GpuIndexIVF::GpuIndexIVF(
20 GpuResourcesProvider* provider,
21 int dims,
22 faiss::MetricType metric,
23 float metricArg,
24 int nlistIn,
25 GpuIndexIVFConfig config)
26 : GpuIndex(provider->getResources(), dims, metric, metricArg, config),
27 nlist(nlistIn),
28 nprobe(1),
29 quantizer(nullptr),
30 ivfConfig_(config) {
31 init_();
32
33 // Only IP and L2 are supported for now
34 if (!(metric_type == faiss::METRIC_L2 ||
35 metric_type == faiss::METRIC_INNER_PRODUCT)) {
36 FAISS_THROW_FMT("unsupported metric type %d", (int)metric_type);
37 }
38 }
39
init_()40 void GpuIndexIVF::init_() {
41 FAISS_THROW_IF_NOT_MSG(nlist > 0, "nlist must be > 0");
42
43 // Spherical by default if the metric is inner_product
44 if (metric_type == faiss::METRIC_INNER_PRODUCT) {
45 cp.spherical = true;
46 }
47
48 // here we set a low # iterations because this is typically used
49 // for large clusterings
50 cp.niter = 10;
51 cp.verbose = verbose;
52
53 if (!quantizer) {
54 // Construct an empty quantizer
55 GpuIndexFlatConfig config = ivfConfig_.flatConfig;
56 // FIXME: inherit our same device
57 config.device = config_.device;
58
59 if (metric_type == faiss::METRIC_L2) {
60 quantizer = new GpuIndexFlatL2(resources_, d, config);
61 } else if (metric_type == faiss::METRIC_INNER_PRODUCT) {
62 quantizer = new GpuIndexFlatIP(resources_, d, config);
63 } else {
64 // unknown metric type
65 FAISS_THROW_FMT("unsupported metric type %d", (int)metric_type);
66 }
67 }
68 }
69
~GpuIndexIVF()70 GpuIndexIVF::~GpuIndexIVF() {
71 delete quantizer;
72 }
73
getQuantizer()74 GpuIndexFlat* GpuIndexIVF::getQuantizer() {
75 return quantizer;
76 }
77
copyFrom(const faiss::IndexIVF * index)78 void GpuIndexIVF::copyFrom(const faiss::IndexIVF* index) {
79 DeviceScope scope(config_.device);
80
81 GpuIndex::copyFrom(index);
82
83 FAISS_ASSERT(index->nlist > 0);
84 FAISS_THROW_IF_NOT_FMT(
85 index->nlist <= (Index::idx_t)std::numeric_limits<int>::max(),
86 "GPU index only supports %zu inverted lists",
87 (size_t)std::numeric_limits<int>::max());
88 nlist = index->nlist;
89
90 FAISS_THROW_IF_NOT_FMT(
91 index->nprobe > 0 && index->nprobe <= getMaxKSelection(),
92 "GPU index only supports nprobe <= %zu; passed %zu",
93 (size_t)getMaxKSelection(),
94 index->nprobe);
95 nprobe = index->nprobe;
96
97 // The metric type may have changed as well, so we might have to
98 // change our quantizer
99 delete quantizer;
100 quantizer = nullptr;
101
102 // Construct an empty quantizer
103 GpuIndexFlatConfig config = ivfConfig_.flatConfig;
104 // FIXME: inherit our same device
105 config.device = config_.device;
106
107 if (index->metric_type == faiss::METRIC_L2) {
108 // FIXME: 2 different float16 options?
109 quantizer = new GpuIndexFlatL2(resources_, this->d, config);
110 } else if (index->metric_type == faiss::METRIC_INNER_PRODUCT) {
111 // FIXME: 2 different float16 options?
112 quantizer = new GpuIndexFlatIP(resources_, this->d, config);
113 } else {
114 // unknown metric type
115 FAISS_ASSERT(false);
116 }
117
118 if (!index->is_trained) {
119 // copied in GpuIndex::copyFrom
120 FAISS_ASSERT(!is_trained && ntotal == 0);
121 return;
122 }
123
124 // copied in GpuIndex::copyFrom
125 // ntotal can exceed max int, but the number of vectors per inverted
126 // list cannot exceed this. We check this in the subclasses.
127 FAISS_ASSERT(is_trained && (ntotal == index->ntotal));
128
129 // Since we're trained, the quantizer must have data
130 FAISS_ASSERT(index->quantizer->ntotal > 0);
131
132 // Right now, we can only handle IndexFlat or derived classes
133 auto qFlat = dynamic_cast<faiss::IndexFlat*>(index->quantizer);
134 FAISS_THROW_IF_NOT_MSG(
135 qFlat,
136 "Only IndexFlat is supported for the coarse quantizer "
137 "for copying from an IndexIVF into a GpuIndexIVF");
138
139 quantizer->copyFrom(qFlat);
140 }
141
copyTo(faiss::IndexIVF * index) const142 void GpuIndexIVF::copyTo(faiss::IndexIVF* index) const {
143 DeviceScope scope(config_.device);
144
145 //
146 // Index information
147 //
148 GpuIndex::copyTo(index);
149
150 //
151 // IndexIVF information
152 //
153 index->nlist = nlist;
154 index->nprobe = nprobe;
155
156 // Construct and copy the appropriate quantizer
157 faiss::IndexFlat* q = nullptr;
158
159 if (this->metric_type == faiss::METRIC_L2) {
160 q = new faiss::IndexFlatL2(this->d);
161
162 } else if (this->metric_type == faiss::METRIC_INNER_PRODUCT) {
163 q = new faiss::IndexFlatIP(this->d);
164
165 } else {
166 // we should have one of the above metrics
167 FAISS_ASSERT(false);
168 }
169
170 FAISS_ASSERT(quantizer);
171 quantizer->copyTo(q);
172
173 if (index->own_fields) {
174 delete index->quantizer;
175 }
176
177 index->quantizer = q;
178 index->quantizer_trains_alone = 0;
179 index->own_fields = true;
180 index->cp = this->cp;
181 index->make_direct_map(false);
182 }
183
getNumLists() const184 int GpuIndexIVF::getNumLists() const {
185 return nlist;
186 }
187
setNumProbes(int nprobe)188 void GpuIndexIVF::setNumProbes(int nprobe) {
189 FAISS_THROW_IF_NOT_FMT(
190 nprobe > 0 && nprobe <= getMaxKSelection(),
191 "GPU index only supports nprobe <= %d; passed %d",
192 getMaxKSelection(),
193 nprobe);
194 this->nprobe = nprobe;
195 }
196
getNumProbes() const197 int GpuIndexIVF::getNumProbes() const {
198 return nprobe;
199 }
200
addImplRequiresIDs_() const201 bool GpuIndexIVF::addImplRequiresIDs_() const {
202 // All IVF indices have storage for IDs
203 return true;
204 }
205
trainQuantizer_(Index::idx_t n,const float * x)206 void GpuIndexIVF::trainQuantizer_(Index::idx_t n, const float* x) {
207 if (n == 0) {
208 // nothing to do
209 return;
210 }
211
212 if (quantizer->is_trained && (quantizer->ntotal == nlist)) {
213 if (this->verbose) {
214 printf("IVF quantizer does not need training.\n");
215 }
216
217 return;
218 }
219
220 if (this->verbose) {
221 printf("Training IVF quantizer on %ld vectors in %dD\n", n, d);
222 }
223
224 DeviceScope scope(config_.device);
225
226 // leverage the CPU-side k-means code, which works for the GPU
227 // flat index as well
228 quantizer->reset();
229 Clustering clus(this->d, nlist, this->cp);
230 clus.verbose = verbose;
231 clus.train(n, x, *quantizer);
232 quantizer->is_trained = true;
233
234 FAISS_ASSERT(quantizer->ntotal == nlist);
235 }
236
237 } // namespace gpu
238 } // namespace faiss
239