All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
GpuIndexIVF.cu
/**
* 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 "GpuIndexIVF.h"
#include "../FaissAssert.h"
#include "../IndexFlat.h"
#include "../IndexIVF.h"
#include "GpuIndexFlat.h"
#include "utils/DeviceUtils.h"
#include "utils/Float16.cuh"
namespace faiss { namespace gpu {
GpuIndexIVF::GpuIndexIVF(GpuResources* resources,
int device,
IndicesOptions indicesOptions,
bool useFloat16CoarseQuantizer,
int dims,
faiss::MetricType metric,
int nlist) :
GpuIndex(resources, device, dims, metric),
indicesOptions_(indicesOptions),
useFloat16CoarseQuantizer_(useFloat16CoarseQuantizer),
nlist_(nlist),
nprobe_(1),
quantizer_(nullptr),
ownsQuantizer_(true) {
#ifndef FAISS_USE_FLOAT16
FAISS_ASSERT(!useFloat16CoarseQuantizer_,
"float16 unsupported; need CUDA SDK >= 7.5");
#endif
init_();
}
GpuIndexIVF::GpuIndexIVF(GpuResources* resources,
int device,
IndicesOptions indicesOptions,
int dims,
faiss::MetricType metric,
int nlist,
GpuIndexFlat* quantizer) :
GpuIndex(resources, device, dims, metric),
indicesOptions_(indicesOptions),
useFloat16CoarseQuantizer_(quantizer->getUseFloat16()),
nlist_(nlist),
nprobe_(1),
quantizer_(quantizer),
ownsQuantizer_(false) {
#ifndef FAISS_USE_FLOAT16
FAISS_ASSERT(!useFloat16CoarseQuantizer_,
"float16 unsupported; need CUDA SDK >= 7.5");
#endif
FAISS_ASSERT(quantizer_->d == this->d);
FAISS_ASSERT(quantizer_->metric_type == this->metric_type);
init_();
}
void
GpuIndexIVF::init_() {
FAISS_ASSERT(nlist_ > 0);
// Spherical by default if the metric is inner_product
if (this->metric_type == faiss::METRIC_INNER_PRODUCT) {
cp_.spherical = true;
}
// here we set a low # iterations because this is typically used
// for large clusterings
cp_.niter = 10;
cp_.verbose = this->verbose;
if (!quantizer_) {
// Construct an empty quantizer
if (this->metric_type == faiss::METRIC_L2) {
// FIXME: 2 different float16 options?
quantizer_ = new GpuIndexFlatL2(resources_, device_, this->d,
useFloat16CoarseQuantizer_);
} else if (this->metric_type == faiss::METRIC_INNER_PRODUCT) {
// FIXME: 2 different float16 options?
quantizer_ = new GpuIndexFlatIP(resources_, device_, this->d,
useFloat16CoarseQuantizer_);
} else {
// unknown metric type
FAISS_ASSERT(false);
}
}
}
GpuIndexIVF::~GpuIndexIVF() {
if (ownsQuantizer_) {
delete quantizer_;
}
}
IndicesOptions
GpuIndexIVF::getIndicesOptions() const {
return indicesOptions_;
}
bool
GpuIndexIVF::getUseFloat16CoarseQuantizer() const {
return useFloat16CoarseQuantizer_;
}
void
GpuIndexIVF::copyFrom(const faiss::IndexIVF* index) {
DeviceScope scope(device_);
this->d = index->d;
this->metric_type = index->metric_type;
FAISS_ASSERT(index->nlist > 0);
FAISS_ASSERT(index->nlist <=
(faiss::Index::idx_t) std::numeric_limits<int>::max());
nlist_ = index->nlist;
nprobe_ = index->nprobe;
// The metric type may have changed as well, so we might have to
// change our quantizer
delete quantizer_;
quantizer_ = nullptr;
if (index->metric_type == faiss::METRIC_L2) {
// FIXME: 2 different float16 options?
quantizer_ = new GpuIndexFlatL2(resources_, device_, this->d,
useFloat16CoarseQuantizer_);
} else if (index->metric_type == faiss::METRIC_INNER_PRODUCT) {
// FIXME: 2 different float16 options?
quantizer_ = new GpuIndexFlatIP(resources_, device_, this->d,
useFloat16CoarseQuantizer_);
} else {
// unknown metric type
FAISS_ASSERT(false);
}
if (!index->is_trained) {
this->is_trained = false;
this->ntotal = 0;
return;
}
// Otherwise, we can populate ourselves from the other index
this->is_trained = true;
// Only use `int` on GPU
FAISS_ASSERT(index->ntotal <=
(faiss::Index::idx_t) std::numeric_limits<int>::max());
this->ntotal = index->ntotal;
// Since we're trained, the quantizer must have data
FAISS_ASSERT(index->quantizer->ntotal > 0);
if (index->metric_type == faiss::METRIC_L2) {
auto q = dynamic_cast< faiss::IndexFlatL2*>(index->quantizer);
FAISS_ASSERT(q);
quantizer_->copyFrom(q);
} else if (index->metric_type == faiss::METRIC_INNER_PRODUCT) {
auto q = dynamic_cast< faiss::IndexFlatIP*>(index->quantizer);
FAISS_ASSERT(q);
quantizer_->copyFrom(q);
} else {
// unknown metric type
FAISS_ASSERT(false);
}
}
void
GpuIndexIVF::copyTo(faiss::IndexIVF* index) const {
DeviceScope scope(device_);
//
// Index information
//
index->ntotal = this->ntotal;
index->d = this->d;
index->metric_type = this->metric_type;
index->is_trained = this->is_trained;
//
// IndexIVF information
//
index->nlist = nlist_;
index->nprobe = nprobe_;
// Construct and copy the appropriate quantizer
faiss::IndexFlat* q = nullptr;
if (this->metric_type == faiss::METRIC_L2) {
q = new faiss::IndexFlatL2(this->d);
} else if (this->metric_type == faiss::METRIC_INNER_PRODUCT) {
q = new faiss::IndexFlatIP(this->d);
} else {
// unknown metric type
FAISS_ASSERT(false);
}
FAISS_ASSERT(quantizer_);
quantizer_->copyTo(q);
if (index->own_fields) {
delete index->quantizer;
}
index->quantizer = q;
index->quantizer_trains_alone = false;
index->own_fields = true;
index->cp = cp_;
index->ids.clear();
index->ids.resize(nlist_);
index->maintain_direct_map = false;
index->direct_map.clear();
}
int
GpuIndexIVF::getNumLists() const {
return nlist_;
}
void
GpuIndexIVF::setNumProbes(int nprobe) {
FAISS_ASSERT(nprobe > 0);
nprobe_ = nprobe;
}
int
GpuIndexIVF::getNumProbes() const {
return nprobe_;
}
void
GpuIndexIVF::add(Index::idx_t n, const float* x) {
// FIXME: GPU-ize
std::vector<Index::idx_t> ids(n);
for (Index::idx_t i = 0; i < n; ++i) {
ids[i] = this->ntotal + i;
}
add_with_ids(n, x, ids.data());
}
void
GpuIndexIVF::trainQuantizer_(faiss::Index::idx_t n, const float* x) {
if (n == 0) {
// nothing to do
return;
}
if (quantizer_->is_trained && (quantizer_->ntotal == nlist_)) {
if (this->verbose) {
printf ("IVF quantizer does not need training.\n");
}
return;
}
if (this->verbose) {
printf ("Training IVF quantizer on %ld vectors in %dD\n", n, d);
}
DeviceScope scope(device_);
// leverage the CPU-side k-means code, which works for the GPU
// flat index as well
quantizer_->reset();
Clustering clus(this->d, nlist_, cp_);
clus.verbose = verbose;
clus.train(n, x, *quantizer_);
quantizer_->is_trained = true;
FAISS_ASSERT(quantizer_->ntotal == nlist_);
}
} } // namespace
faiss::IndexIVF
Definition: IndexIVF.h:46
faiss::gpu::GpuIndexIVF::getIndicesOptions
IndicesOptions getIndicesOptions() const
What indices storage options are we using?
Definition: GpuIndexIVF.cu:107
faiss::gpu::GpuIndexIVF::getNumProbes
int getNumProbes() const
Returns our current number of list probes per query.
Definition: GpuIndexIVF.cu:241
faiss::IndexFlat
Definition: IndexFlat.h:24
faiss::gpu::GpuIndexIVF::setNumProbes
void setNumProbes(int nprobe)
Sets the number of list probes per query.
Definition: GpuIndexIVF.cu:235
faiss::ClusteringParameters::niter
int niter
clustering iterations
Definition: Clustering.h:26
faiss::gpu::GpuIndexIVF::useFloat16CoarseQuantizer_
const bool useFloat16CoarseQuantizer_
Definition: GpuIndexIVF.h:85
faiss::gpu::GpuIndexIVF::getNumLists
int getNumLists() const
Returns the number of inverted lists we&#39;re managing.
Definition: GpuIndexIVF.cu:230
faiss::gpu::GpuIndexFlat::copyTo
void copyTo(faiss::IndexFlat *index) const
Definition: GpuIndexFlat.cu:107
faiss::IndexIVF::nprobe
size_t nprobe
number of probes at query time
Definition: IndexIVF.h:48
faiss::gpu::GpuIndexIVF::getUseFloat16CoarseQuantizer
bool getUseFloat16CoarseQuantizer() const
Is our coarse quantizer storing and performing math in float16?
Definition: GpuIndexIVF.cu:112
faiss::IndexIVF::quantizer_trains_alone
bool quantizer_trains_alone
just pass over the trainset to quantizer
Definition: IndexIVF.h:51
faiss::IndexFlatL2
Definition: IndexFlat.h:81
faiss::Index::add_with_ids
virtual void add_with_ids(idx_t n, const float *x, const long *xids)
Definition: Index.cpp:32
faiss::IndexFlatIP
Definition: IndexFlat.h:75
faiss::gpu::GpuIndex::device_
int device_
The GPU device we are resident on.
Definition: GpuIndex.h:43
faiss::IndexIVF::ids
std::vector< std::vector< long > > ids
Inverted lists for indexes.
Definition: IndexIVF.h:56
faiss::Index::d
int d
vector dimension
Definition: Index.h:66
faiss::IndexIVF::quantizer
Index * quantizer
quantizer that maps vectors to inverted lists
Definition: IndexIVF.h:50
faiss::IndexIVF::cp
ClusteringParameters cp
to override default clustering params
Definition: IndexIVF.h:54
faiss::gpu::GpuIndexIVF::nprobe_
int nprobe_
Number of inverted list probes per query.
Definition: GpuIndexIVF.h:91
faiss::IndexIVF::own_fields
bool own_fields
whether object owns the quantizer
Definition: IndexIVF.h:52
faiss::gpu::GpuIndexIVF::indicesOptions_
const IndicesOptions indicesOptions_
How should indices be stored on the GPU?
Definition: GpuIndexIVF.h:81
faiss::gpu::GpuIndex::resources_
GpuResources * resources_
Manages streans, cuBLAS handles and scratch memory for devices.
Definition: GpuIndex.h:40
faiss::gpu::GpuIndexIVF::copyTo
void copyTo(faiss::IndexIVF *index) const
Copy what we have to the CPU equivalent.
Definition: GpuIndexIVF.cu:181
faiss::Index::idx_t
long idx_t
all indices are this type
Definition: Index.h:64
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:67
faiss::Index::verbose
bool verbose
verbosity level
Definition: Index.h:68
faiss::gpu::GpuIndexFlatIP
Definition: GpuIndexFlat.h:158
faiss::gpu::GpuIndexFlat::copyFrom
void copyFrom(const faiss::IndexFlat *index)
Definition: GpuIndexFlat.cu:81
faiss::gpu::DeviceScope
Definition: DeviceUtils.h:55
faiss::gpu::GpuIndexIVF::quantizer_
GpuIndexFlat * quantizer_
Quantizer for inverted lists.
Definition: GpuIndexIVF.h:97
faiss::gpu::GpuIndexFlatL2
Definition: GpuIndexFlat.h:131
faiss::Index::metric_type
MetricType metric_type
type of metric this index uses for search
Definition: Index.h:74
faiss::IndexIVF::nlist
size_t nlist
number of possible key values
Definition: IndexIVF.h:47
faiss::gpu::GpuIndexIVF::cp_
ClusteringParameters cp_
Ability to override default clustering parameters.
Definition: GpuIndexIVF.h:94
faiss::gpu::GpuIndexFlat::reset
void reset() override
Clears all vectors from this index.
Definition: GpuIndexFlat.cu:135
faiss::gpu::GpuIndexIVF::copyFrom
void copyFrom(const faiss::IndexIVF *index)
Copy what we need from the CPU equivalent.
Definition: GpuIndexIVF.cu:117
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::gpu::GpuIndexIVF::ownsQuantizer_
bool ownsQuantizer_
Do we own the above quantizer instance?
Definition: GpuIndexIVF.h:100
faiss::gpu::GpuIndexIVF::add
void add(Index::idx_t n, const float *x) override
Definition: GpuIndexIVF.cu:247
faiss::IndexIVF::maintain_direct_map
bool maintain_direct_map
map for direct access to the elements. Enables reconstruct().
Definition: IndexIVF.h:59
faiss::ClusteringParameters::spherical
bool spherical
do we want normalized centroids?
Definition: Clustering.h:31
faiss::MetricType
MetricType
Some algorithms support both an inner product vetsion and a L2 search version.
Definition: Index.h:44