All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
IndexProxy.cpp
/**
* Copyright (c) 2015-present, Facebook, Inc.
* All rights reserved.
*
* This source code is licensed under the CC-by-NC license found in the
* LICENSE file in the root directory of this source tree.
*/
// Copyright 2004-present Facebook. All Rights Reserved.
#include "IndexProxy.h"
#include "../FaissAssert.h"
#include "../Clustering.h"
#include "GpuIndexFlat.h"
#include "StandardGpuResources.h"
#include <cstring>
namespace faiss { namespace gpu {
IndexProxy::IndexProxy():own_fields(false) {
}
IndexProxy::~IndexProxy() {
if (own_fields) {
for (auto& index : indices_)
delete index.first;
}
}
void
IndexProxy::addIndex(faiss::Index* index) {
// Make sure that the parameters are the same for all prior indices
if (!indices_.empty()) {
auto& existing = indices_.front().first;
if (index->d != existing->d) {
FAISS_ASSERT(false);
return;
}
if (index->ntotal != existing->ntotal) {
FAISS_ASSERT(false);
return;
}
if (index->metric_type != existing->metric_type) {
FAISS_ASSERT(false);
return;
}
} else {
// Set our parameters
// FIXME: this is a little bit weird
this->d = index->d;
this->ntotal = index->ntotal;
this->verbose = index->verbose;
this->is_trained = index->is_trained;
this->metric_type = index->metric_type;
}
indices_.emplace_back(
std::make_pair(index,
std::unique_ptr<WorkerThread>(new WorkerThread)));
}
void
IndexProxy::removeIndex(faiss::Index* index) {
for (auto it = indices_.begin(); it != indices_.end(); ++it) {
if (it->first == index) {
// This is our index; stop the worker thread before removing it,
// to ensure that it has finished before function exit
it->second->stop();
it->second->waitForThreadExit();
indices_.erase(it);
return;
}
}
// index not found
FAISS_ASSERT(false);
}
void
IndexProxy::runOnIndex(std::function<void(faiss::Index*)> f) {
std::vector<std::future<bool>> v;
for (auto& index : indices_) {
auto indexPtr = index.first;
v.emplace_back(index.second->add([indexPtr, f](){ f(indexPtr); }));
}
// Blocking wait for completion
for (auto& func : v) {
func.get();
}
}
void
IndexProxy::reset() {
runOnIndex([](faiss::Index* index){ index->reset(); });
ntotal = 0;
}
void
IndexProxy::train(Index::idx_t n, const float* x) {
runOnIndex([n, x](faiss::Index* index){ index->train(n, x); });
}
void
IndexProxy::add(Index::idx_t n, const float* x) {
runOnIndex([n, x](faiss::Index* index){ index->add(n, x); });
ntotal += n;
}
void
IndexProxy::reconstruct(Index::idx_t n, float* x) const {
FAISS_ASSERT (count() > 0);
indices_[0].first->reconstruct (n, x);
}
void
IndexProxy::search(faiss::Index::idx_t n,
const float* x,
faiss::Index::idx_t k,
float* distances,
faiss::Index::idx_t* labels) const {
FAISS_ASSERT(!indices_.empty());
if (indices_.empty()) {
return;
}
auto dim = indices_.front().first->d;
std::vector<std::future<bool>> v;
// Partition the query by the number of indices we have
auto queriesPerIndex =
(faiss::Index::idx_t) (n + indices_.size() - 1) / indices_.size();
FAISS_ASSERT(n / queriesPerIndex <= indices_.size());
for (int i = 0; i < indices_.size(); ++i) {
auto base = i * queriesPerIndex;
if (base >= n) {
break;
}
auto numForIndex = std::min(queriesPerIndex, n - base);
auto queryStart = x + base * dim;
auto distancesStart = distances + base * k;
auto labelsStart = labels + base * k;
auto indexPtr = indices_[i].first;
auto fn =
[indexPtr, numForIndex, queryStart, k, distancesStart, labelsStart]() {
indexPtr->search(numForIndex, queryStart,
k, distancesStart, labelsStart);
};
v.emplace_back(indices_[i].second->add(std::move(fn)));
}
// Blocking wait for completion
for (auto& f : v) {
f.get();
}
}
void
IndexProxy::set_typename() {
// FIXME: implement
FAISS_ASSERT(false);
}
//
// GPU clustering implementation
//
float kmeans_clustering_gpu (int ngpu, size_t d, size_t n, size_t k,
const float *x,
float *centroids,
bool useFloat16,
bool storeTransposed)
{
Clustering clus (d, k);
// display logs if > 16Gflop per iteration
clus.verbose = d * n * k > (1L << 34);
FAISS_ASSERT(ngpu >= 1);
std::vector<std::unique_ptr<StandardGpuResources> > res;
std::vector<std::unique_ptr<GpuIndexFlatL2> > sub_indices;
for(int dev_no = 0; dev_no < ngpu; dev_no++) {
res.emplace_back(new StandardGpuResources());
GpuIndexFlatConfig config;
config.device = dev_no;
config.useFloat16 = useFloat16;
config.storeTransposed = storeTransposed;
sub_indices.emplace_back(
new GpuIndexFlatL2(res.back().get(), d, config));
}
IndexProxy proxy;
Index *index;
if (ngpu == 1) {
index = sub_indices[0].get();
} else {
for(int dev_no = 0; dev_no < ngpu; dev_no++) {
proxy.addIndex(sub_indices[dev_no].get());
}
index = &proxy;
}
clus.train (n, x, *index);
memcpy(centroids, clus.centroids.data(), sizeof(*centroids) * d * k);
return clus.obj.back();
}
} } // 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