Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

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