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