Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
IndexProxy.cpp
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 "IndexProxy.h"
11 #include "../FaissAssert.h"
12
13 #include "../Clustering.h"
14 #include "GpuIndexFlat.h"
15 #include "StandardGpuResources.h"
16 #include <cstring>
17
18 namespace faiss { namespace gpu {
19
20 IndexProxy::IndexProxy():own_fields(false) {
21 }
22
23 IndexProxy::~IndexProxy() {
24 if (own_fields) {
25 for (auto& index : indices_)
26 delete index.first;
27 }
28 }
29
30 void
31 IndexProxy::addIndex(faiss::Index* index) {
32 // Make sure that the parameters are the same for all prior indices
33 if (!indices_.empty()) {
34 auto& existing = indices_.front().first;
35
36 if (index->d != existing->d) {
37 FAISS_ASSERT(false);
38 return;
39 }
40
41 if (index->ntotal != existing->ntotal) {
42 FAISS_ASSERT(false);
43 return;
44 }
45
46 if (index->metric_type != existing->metric_type) {
47 FAISS_ASSERT(false);
48 return;
49 }
50 } else {
51 // Set our parameters
52 // FIXME: this is a little bit weird
53 this->d = index->d;
54 this->ntotal = index->ntotal;
55 this->verbose = index->verbose;
56 this->is_trained = index->is_trained;
57 this->metric_type = index->metric_type;
58 }
59
60 indices_.emplace_back(
61 std::make_pair(index,
62 std::unique_ptr<WorkerThread>(new WorkerThread)));
63 }
64
65 void
66 IndexProxy::removeIndex(faiss::Index* index) {
67 for (auto it = indices_.begin(); it != indices_.end(); ++it) {
68 if (it->first == index) {
69 // This is our index; stop the worker thread before removing it,
70 // to ensure that it has finished before function exit
71 it->second->stop();
72 it->second->waitForThreadExit();
73
74 indices_.erase(it);
75 return;
76 }
77 }
78
79 // index not found
80 FAISS_ASSERT(false);
81 }
82
83 void
84 IndexProxy::runOnIndex(std::function<void(faiss::Index*)> f) {
85 std::vector<std::future<bool>> v;
86
87 for (auto& index : indices_) {
88 auto indexPtr = index.first;
89 v.emplace_back(index.second->add([indexPtr, f](){ f(indexPtr); }));
90 }
91
92 // Blocking wait for completion
93 for (auto& func : v) {
94 func.get();
95 }
96 }
97
98 void
99 IndexProxy::reset() {
100 runOnIndex([](faiss::Index* index){ index->reset(); });
101 ntotal = 0;
102 }
103
104 void
105 IndexProxy::train(Index::idx_t n, const float* x) {
106 runOnIndex([n, x](faiss::Index* index){ index->train(n, x); });
107 }
108
109 void
110 IndexProxy::add(Index::idx_t n, const float* x) {
111 runOnIndex([n, x](faiss::Index* index){ index->add(n, x); });
112 ntotal += n;
113 }
114
115 void
116 IndexProxy::reconstruct(Index::idx_t n, float* x) const {
117 FAISS_ASSERT (count() > 0);
118 indices_[0].first->reconstruct (n, x);
119 }
120
121
122 void
123 IndexProxy::search(faiss::Index::idx_t n,
124 const float* x,
125 faiss::Index::idx_t k,
126 float* distances,
127 faiss::Index::idx_t* labels) const {
128 FAISS_ASSERT(!indices_.empty());
129 if (n == 0) {
130 return;
131 }
132
133 auto dim = indices_.front().first->d;
134
135 std::vector<std::future<bool>> v;
136
137 // Partition the query by the number of indices we have
138 auto queriesPerIndex =
139 (faiss::Index::idx_t) (n + indices_.size() - 1) / indices_.size();
140 FAISS_ASSERT(n / queriesPerIndex <= indices_.size());
141
142 for (int i = 0; i < indices_.size(); ++i) {
143 auto base = i * queriesPerIndex;
144 if (base >= n) {
145 break;
146 }
147
148 auto numForIndex = std::min(queriesPerIndex, n - base);
149 auto queryStart = x + base * dim;
150 auto distancesStart = distances + base * k;
151 auto labelsStart = labels + base * k;
152
153 auto indexPtr = indices_[i].first;
154 auto fn =
155 [indexPtr, numForIndex, queryStart, k, distancesStart, labelsStart]() {
156 indexPtr->search(numForIndex, queryStart,
157 k, distancesStart, labelsStart);
158 };
159
160 v.emplace_back(indices_[i].second->add(std::move(fn)));
161 }
162
163 // Blocking wait for completion
164 for (auto& f : v) {
165 f.get();
166 }
167 }
168
169
170
171 //
172 // GPU clustering implementation
173 //
174
175 float kmeans_clustering_gpu (int ngpu, size_t d, size_t n, size_t k,
176 const float *x,
177 float *centroids,
178 bool useFloat16,
179 bool storeTransposed)
180 {
181 Clustering clus (d, k);
182 // display logs if > 16Gflop per iteration
183 clus.verbose = d * n * k > (1L << 34);
184 FAISS_ASSERT(ngpu >= 1);
185
186 std::vector<std::unique_ptr<StandardGpuResources> > res;
187 std::vector<std::unique_ptr<GpuIndexFlatL2> > sub_indices;
188 for(int dev_no = 0; dev_no < ngpu; dev_no++) {
189 res.emplace_back(new StandardGpuResources());
190
191
192 GpuIndexFlatConfig config;
193 config.device = dev_no;
194 config.useFloat16 = useFloat16;
195 config.storeTransposed = storeTransposed;
196
197 sub_indices.emplace_back(
198 new GpuIndexFlatL2(res.back().get(), d, config));
199 }
200
201 IndexProxy proxy;
202 Index *index;
203 if (ngpu == 1) {
204 index = sub_indices[0].get();
205 } else {
206 for(int dev_no = 0; dev_no < ngpu; dev_no++) {
207 proxy.addIndex(sub_indices[dev_no].get());
208 }
209 index = &proxy;
210 }
211 clus.train (n, x, *index);
212
213 memcpy(centroids, clus.centroids.data(), sizeof(*centroids) * d * k);
214 return clus.obj.back();
215
216 }
217
218
219
220
221 } } // namespace
faiss::gpu::GpuIndexFlatConfig::storeTransposed
bool storeTransposed
Definition: GpuIndexFlat.h:47
faiss::gpu::IndexProxy::removeIndex
void removeIndex(faiss::Index *index)
Definition: IndexProxy.cpp:66
faiss::gpu::IndexProxy::runOnIndex
void runOnIndex(std::function< void(faiss::Index *)> f)
Definition: IndexProxy.cpp:84
faiss::Index::reset
virtual void reset()=0
removes all elements from the database.
faiss::gpu::WorkerThread
Definition: WorkerThread.h:19
faiss::Index::train
virtual void train(idx_t n, const float *x)
Definition: Index.cpp:24
faiss::gpu::GpuIndexFlatConfig::useFloat16
bool useFloat16
Whether or not data is stored as float16.
Definition: GpuIndexFlat.h:34
faiss::gpu::GpuIndexConfig::device
int device
GPU device on which the index is resident.
Definition: GpuIndex.h:26
faiss::Index::d
int d
vector dimension
Definition: Index.h:66
faiss::gpu::GpuIndexFlatConfig
Definition: GpuIndexFlat.h:26
faiss::Clustering
Definition: Clustering.h:56
faiss::Index::add
virtual void add(idx_t n, const float *x)=0
faiss::gpu::StandardGpuResources
Definition: StandardGpuResources.h:22
faiss::Index::idx_t
long idx_t
all indices are this type
Definition: Index.h:64
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::IndexProxy::addIndex
void addIndex(faiss::Index *index)
Definition: IndexProxy.cpp:31
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::Index
Definition: Index.h:62
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