Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

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