All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
GpuIndexIVFFlat.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 "GpuIndexIVFFlat.h"
#include "../IndexFlat.h"
#include "../IndexIVF.h"
#include "GpuIndexFlat.h"
#include "GpuResources.h"
#include "impl/IVFFlat.cuh"
#include "utils/CopyUtils.cuh"
#include "utils/DeviceUtils.h"
#include "utils/Float16.cuh"
#include <limits>
namespace faiss { namespace gpu {
GpuIndexIVFFlat::GpuIndexIVFFlat(GpuResources* resources,
const faiss::IndexIVFFlat* index,
GpuIndexIVFFlatConfig config) :
GpuIndexIVF(resources,
index->d,
index->metric_type,
index->nlist,
config),
ivfFlatConfig_(config),
reserveMemoryVecs_(0),
index_(nullptr) {
#ifndef FAISS_USE_FLOAT16
FAISS_THROW_IF_NOT_MSG(!ivfFlatConfig_.useFloat16IVFStorage,
"float16 unsupported; need CUDA SDK >= 7.5");
#endif
copyFrom(index);
}
GpuIndexIVFFlat::GpuIndexIVFFlat(GpuResources* resources,
int dims,
int nlist,
faiss::MetricType metric,
GpuIndexIVFFlatConfig config) :
GpuIndexIVF(resources, dims, metric, nlist, config),
ivfFlatConfig_(config),
reserveMemoryVecs_(0),
index_(nullptr) {
// faiss::Index params
this->is_trained = false;
#ifndef FAISS_USE_FLOAT16
FAISS_THROW_IF_NOT_MSG(!ivfFlatConfig_.useFloat16IVFStorage,
"float16 unsupported; need CUDA SDK >= 7.5");
#endif
// We haven't trained ourselves, so don't construct the IVFFlat
// index yet
}
GpuIndexIVFFlat::~GpuIndexIVFFlat() {
delete index_;
}
void
GpuIndexIVFFlat::reserveMemory(size_t numVecs) {
reserveMemoryVecs_ = numVecs;
if (index_) {
index_->reserveMemory(numVecs);
}
}
void
GpuIndexIVFFlat::copyFrom(const faiss::IndexIVFFlat* index) {
DeviceScope scope(device_);
GpuIndexIVF::copyFrom(index);
// Clear out our old data
delete index_;
index_ = nullptr;
// The other index might not be trained
if (!index->is_trained) {
return;
}
// Otherwise, we can populate ourselves from the other index
this->is_trained = true;
// Copy our lists as well
index_ = new IVFFlat(resources_,
quantizer_->getGpuData(),
index->metric_type == faiss::METRIC_L2,
ivfFlatConfig_.useFloat16IVFStorage,
ivfFlatConfig_.indicesOptions,
memorySpace_);
FAISS_ASSERT(index->vecs.size() == index->ids.size());
for (size_t i = 0; i < index->vecs.size(); ++i) {
auto& vecs = index->vecs[i];
auto& ids = index->ids[i];
FAISS_ASSERT(vecs.size() % this->d == 0);
auto numVecs = vecs.size() / this->d;
FAISS_ASSERT(numVecs == ids.size());
// GPU index can only support max int entries per list
FAISS_THROW_IF_NOT_FMT(ids.size() <=
(size_t) std::numeric_limits<int>::max(),
"GPU inverted list can only support "
"%zu entries; %zu found",
(size_t) std::numeric_limits<int>::max(),
ids.size());
index_->addCodeVectorsFromCpu(i, vecs.data(), ids.data(), numVecs);
}
}
void
GpuIndexIVFFlat::copyTo(faiss::IndexIVFFlat* index) const {
DeviceScope scope(device_);
// We must have the indices in order to copy to ourselves
FAISS_THROW_IF_NOT_MSG(ivfFlatConfig_.indicesOptions != INDICES_IVF,
"Cannot copy to CPU as GPU index doesn't retain "
"indices (INDICES_IVF)");
GpuIndexIVF::copyTo(index);
// Clear out the old inverted lists
index->vecs.clear();
index->vecs.resize(nlist_);
// Copy the inverted lists
if (index_) {
for (int i = 0; i < nlist_; ++i) {
index->vecs[i] = index_->getListVectors(i);
index->ids[i] = index_->getListIndices(i);
}
}
}
size_t
GpuIndexIVFFlat::reclaimMemory() {
if (index_) {
DeviceScope scope(device_);
return index_->reclaimMemory();
}
return 0;
}
void
GpuIndexIVFFlat::reset() {
if (index_) {
DeviceScope scope(device_);
index_->reset();
this->ntotal = 0;
} else {
FAISS_ASSERT(this->ntotal == 0);
}
}
void
GpuIndexIVFFlat::train(Index::idx_t n, const float* x) {
DeviceScope scope(device_);
if (this->is_trained) {
FAISS_ASSERT(quantizer_->is_trained);
FAISS_ASSERT(quantizer_->ntotal == nlist_);
FAISS_ASSERT(index_);
return;
}
FAISS_ASSERT(!index_);
trainQuantizer_(n, x);
// The quantizer is now trained; construct the IVF index
index_ = new IVFFlat(resources_,
quantizer_->getGpuData(),
this->metric_type == faiss::METRIC_L2,
ivfFlatConfig_.useFloat16IVFStorage,
ivfFlatConfig_.indicesOptions,
memorySpace_);
if (reserveMemoryVecs_) {
index_->reserveMemory(reserveMemoryVecs_);
}
this->is_trained = true;
}
void
GpuIndexIVFFlat::addImpl_(Index::idx_t n,
const float* x,
const Index::idx_t* xids) {
// Device is already set in GpuIndex::addInternal_
FAISS_ASSERT(index_);
FAISS_ASSERT(n > 0);
auto stream = resources_->getDefaultStreamCurrentDevice();
auto deviceVecs =
toDevice<float, 2>(resources_,
device_,
const_cast<float*>(x),
stream,
{(int) n, index_->getDim()});
static_assert(sizeof(long) == sizeof(Index::idx_t), "size mismatch");
auto deviceIds =
toDevice<long, 1>(resources_,
device_,
const_cast<long*>(xids),
stream,
{(int) n});
// Not all vectors may be able to be added (some may contain NaNs
// etc)
ntotal += index_->classifyAndAddVectors(deviceVecs, deviceIds);
}
void
GpuIndexIVFFlat::searchImpl_(faiss::Index::idx_t n,
const float* x,
faiss::Index::idx_t k,
float* distances,
faiss::Index::idx_t* labels) const {
// Device is already set in GpuIndex::search
FAISS_ASSERT(index_);
FAISS_ASSERT(n > 0);
auto stream = resources_->getDefaultStream(device_);
// Make sure arguments are on the device we desire; use temporary
// memory allocations to move it if necessary
auto devX =
toDevice<float, 2>(resources_,
device_,
const_cast<float*>(x),
stream,
{(int) n, this->d});
auto devDistances =
toDevice<float, 2>(resources_,
device_,
distances,
stream,
{(int) n, (int) k});
auto devLabels =
toDevice<faiss::Index::idx_t, 2>(resources_,
device_,
labels,
stream,
{(int) n, (int) k});
index_->query(devX, nprobe_, k, devDistances, devLabels);
// Copy back if necessary
fromDevice<float, 2>(devDistances, distances, stream);
fromDevice<faiss::Index::idx_t, 2>(devLabels, labels, stream);
}
} } // namespace
faiss::gpu::IVFFlat
Definition: IVFFlat.cuh:17
faiss::gpu::GpuIndexIVFFlat::GpuIndexIVFFlat
GpuIndexIVFFlat(GpuResources *resources, const faiss::IndexIVFFlat *index, GpuIndexIVFFlatConfig config=GpuIndexIVFFlatConfig())
Definition: GpuIndexIVFFlat.cu:25
faiss::gpu::GpuResources::getDefaultStreamCurrentDevice
cudaStream_t getDefaultStreamCurrentDevice()
Calls getDefaultStream with the current device.
Definition: GpuResources.cpp:25
faiss::gpu::GpuIndexIVFFlat::searchImpl_
void searchImpl_(faiss::Index::idx_t n, const float *x, faiss::Index::idx_t k, float *distances, faiss::Index::idx_t *labels) const override
Called from GpuIndex for search.
Definition: GpuIndexIVFFlat.cu:233
faiss::gpu::IVFBase::getDim
int getDim() const
Return the number of dimensions we are indexing.
Definition: IVFBase.cu:100
faiss::gpu::GpuIndexFlat::getGpuData
FlatIndex * getGpuData()
For internal access.
Definition: GpuIndexFlat.h:120
faiss::gpu::GpuIndexIVFFlat::copyFrom
void copyFrom(const faiss::IndexIVFFlat *index)
Definition: GpuIndexIVFFlat.cu:79
faiss::gpu::IVFBase::reserveMemory
void reserveMemory(size_t numVecs)
Reserve GPU memory in our inverted lists for this number of vectors.
Definition: IVFBase.cu:45
faiss::gpu::GpuIndexIVFFlat::reclaimMemory
size_t reclaimMemory()
Definition: GpuIndexIVFFlat.cu:150
faiss::gpu::GpuIndexIVFFlatConfig::useFloat16IVFStorage
bool useFloat16IVFStorage
Definition: GpuIndexIVFFlat.h:29
faiss::gpu::GpuResources
Definition: GpuResources.h:23
faiss::IndexIVFFlat
Definition: IndexIVF.h:132
faiss::gpu::GpuIndexIVFFlat::addImpl_
void addImpl_(faiss::Index::idx_t n, const float *x, const faiss::Index::idx_t *ids) override
Called from GpuIndex for add/add_with_ids.
Definition: GpuIndexIVFFlat.cu:203
faiss::gpu::GpuIndexIVFFlat::train
void train(Index::idx_t n, const float *x) override
Definition: GpuIndexIVFFlat.cu:173
faiss::IndexIVF::ids
std::vector< std::vector< long > > ids
Inverted lists for indexes.
Definition: IndexIVF.h:55
faiss::gpu::GpuIndexIVF
Definition: GpuIndexIVF.h:37
faiss::Index::d
int d
vector dimension
Definition: Index.h:64
faiss::gpu::IVFFlat::classifyAndAddVectors
int classifyAndAddVectors(Tensor< float, 2, true > &vecs, Tensor< long, 1, true > &indices)
Definition: IVFFlat.cu:130
faiss::gpu::GpuIndexIVFFlat::copyTo
void copyTo(faiss::IndexIVFFlat *index) const
Definition: GpuIndexIVFFlat.cu:126
faiss::gpu::GpuIndexIVF::nprobe_
int nprobe_
Number of inverted list probes per query.
Definition: GpuIndexIVF.h:86
faiss::gpu::GpuIndexIVFFlat::reserveMemory
void reserveMemory(size_t numVecs)
Reserve GPU memory in our inverted lists for this number of vectors.
Definition: GpuIndexIVFFlat.cu:71
faiss::gpu::GpuIndex::device_
const int device_
The GPU device we are resident on.
Definition: GpuIndex.h:94
faiss::gpu::GpuIndex::resources_
GpuResources * resources_
Manages streans, cuBLAS handles and scratch memory for devices.
Definition: GpuIndex.h:91
faiss::gpu::GpuIndexIVF::copyTo
void copyTo(faiss::IndexIVF *index) const
Copy what we have to the CPU equivalent.
Definition: GpuIndexIVF.cu:148
faiss::Index::idx_t
long idx_t
all indices are this type
Definition: Index.h:62
faiss::gpu::GpuIndexIVF::nlist_
int nlist_
Number of inverted lists that we manage.
Definition: GpuIndexIVF.h:83
faiss::Index::ntotal
idx_t ntotal
total nb of indexed vectors
Definition: Index.h:65
faiss::gpu::IVFBase::reset
void reset()
Definition: IVFBase.cu:75
faiss::gpu::IVFFlat::addCodeVectorsFromCpu
void addCodeVectorsFromCpu(int listId, const float *vecs, const long *indices, size_t numVecs)
Definition: IVFFlat.cu:58
faiss::gpu::DeviceScope
Definition: DeviceUtils.h:62
faiss::gpu::GpuIndex::memorySpace_
const MemorySpace memorySpace_
The memory space of our primary storage on the GPU.
Definition: GpuIndex.h:97
faiss::gpu::GpuIndexIVF::quantizer_
GpuIndexFlat * quantizer_
Quantizer for inverted lists.
Definition: GpuIndexIVF.h:92
faiss::gpu::IVFFlat::query
void query(Tensor< float, 2, true > &queries, int nprobe, int k, Tensor< float, 2, true > &outDistances, Tensor< long, 2, true > &outIndices)
Definition: IVFFlat.cu:288
faiss::gpu::IVFFlat::getListVectors
std::vector< float > getListVectors(int listId) const
Return the vectors of a particular list back to the CPU.
Definition: IVFFlat.cu:354
faiss::Index::metric_type
MetricType metric_type
type of metric this index uses for search
Definition: Index.h:72
faiss::gpu::IVFBase::getListIndices
std::vector< long > getListIndices(int listId) const
Return the list indices of a particular list back to the CPU.
Definition: IVFBase.cu:207
faiss::gpu::GpuIndexIVF::copyFrom
void copyFrom(const faiss::IndexIVF *index)
Copy what we need from the CPU equivalent.
Definition: GpuIndexIVF.cu:80
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::gpu::GpuIndexIVFFlatConfig
Definition: GpuIndexIVFFlat.h:22
faiss::gpu::GpuIndexIVFConfig::indicesOptions
IndicesOptions indicesOptions
Index storage options for the GPU.
Definition: GpuIndexIVF.h:31
faiss::gpu::IVFBase::reclaimMemory
size_t reclaimMemory()
Definition: IVFBase.cu:105
faiss::gpu::GpuIndexIVFFlat::reset
void reset() override
removes all elements from the database.
Definition: GpuIndexIVFFlat.cu:161
faiss::MetricType
MetricType
Some algorithms support both an inner product vetsion and a L2 search version.
Definition: Index.h:43
faiss::IndexIVFFlat::vecs
std::vector< std::vector< float > > vecs
Definition: IndexIVF.h:135