9 #include "GpuIndexIVFPQ.h"
10 #include "../IndexFlat.h"
11 #include "../IndexIVFPQ.h"
12 #include "../ProductQuantizer.h"
13 #include "GpuIndexFlat.h"
14 #include "GpuResources.h"
15 #include "impl/IVFPQ.cuh"
16 #include "utils/CopyUtils.cuh"
17 #include "utils/DeviceUtils.h"
21 namespace faiss {
namespace gpu {
34 reserveMemoryVecs_(0),
36 #ifndef FAISS_USE_FLOAT16
56 subQuantizers_(subQuantizers),
57 bitsPerCode_(bitsPerCode),
58 reserveMemoryVecs_(0),
60 #ifndef FAISS_USE_FLOAT16
67 FAISS_ASSERT(this->
metric_type == faiss::METRIC_L2);
73 GpuIndexIVFPQ::~GpuIndexIVFPQ() {
82 FAISS_THROW_IF_NOT_MSG(index->
metric_type == faiss::METRIC_L2,
83 "inner product unsupported");
90 subQuantizers_ = index->
pq.
M;
94 FAISS_ASSERT(index->
pq.
nbits == 8);
124 size_t nlist = ivf ? ivf->
nlist : 0;
125 for (
size_t i = 0; i < nlist; ++i) {
129 FAISS_THROW_IF_NOT_FMT(list_size <=
130 (
size_t) std::numeric_limits<int>::max(),
131 "GPU inverted list can only support "
132 "%zu entries; %zu found",
133 (
size_t) std::numeric_limits<int>::max(),
136 index_->addCodeVectorsFromCpu(
146 FAISS_THROW_IF_NOT_MSG(ivfpqConfig_.
indicesOptions != INDICES_IVF,
147 "Cannot copy to CPU as GPU index doesn't retain "
148 "indices (INDICES_IVF)");
175 for (
int i = 0; i <
nlist_; ++i) {
178 index->
invlists->add_entries (i, ids.size(), ids.data(), codes.data());
183 index->
pq.
centroids.resize(devPQCentroids.numElements());
185 fromDevice<float, 3>(devPQCentroids,
197 reserveMemoryVecs_ = numVecs;
222 return subQuantizers_;
232 return utils::pow2(bitsPerCode_);
253 FAISS_ASSERT(this->
ntotal == 0);
258 GpuIndexIVFPQ::trainResidualQuantizer_(
Index::idx_t n,
const float* x) {
261 n = std::min(n, (
Index::idx_t) (1 << bitsPerCode_) * 64);
264 printf(
"computing residuals\n");
267 std::vector<Index::idx_t>
assign(n);
270 std::vector<float> residuals(n *
d);
272 for (idx_t i = 0; i < n; i++) {
277 printf(
"training %d x %d product quantizer on %ld vectors in %dD\n",
284 pq.train(n, residuals.data());
294 if (reserveMemoryVecs_) {
308 FAISS_ASSERT(index_);
312 FAISS_ASSERT(!index_);
314 trainQuantizer_(n, x);
315 trainResidualQuantizer_(n, x);
317 FAISS_ASSERT(index_);
327 FAISS_ASSERT(index_);
333 static_assert(
sizeof(
long) ==
sizeof(
Index::idx_t),
"size mismatch");
351 FAISS_ASSERT(index_);
358 static_assert(
sizeof(
long) ==
sizeof(
Index::idx_t),
"size mismatch");
361 index_->
query(queries,
nprobe_, k, outDistances, outLabels);
366 FAISS_ASSERT(index_);
370 std::vector<unsigned char>
372 FAISS_ASSERT(index_);
380 FAISS_ASSERT(index_);
387 GpuIndexIVFPQ::verifySettings_()
const {
391 FAISS_THROW_IF_NOT_MSG(
nlist_ > 0,
"nlist must be >0");
394 FAISS_THROW_IF_NOT_FMT(bitsPerCode_ <= 8,
395 "Bits per code must be <= 8 (passed %d)", bitsPerCode_);
398 FAISS_THROW_IF_NOT_FMT(this->
d % subQuantizers_ == 0,
399 "Number of sub-quantizers (%d) must be an "
400 "even divisor of the number of dimensions (%d)",
401 subQuantizers_, this->
d);
405 "Number of bytes per encoded vector / sub-quantizers (%d) "
411 int lookupTableSize =
sizeof(float);
412 #ifdef FAISS_USE_FLOAT16
414 lookupTableSize =
sizeof(half);
420 size_t requiredSmemSize =
421 lookupTableSize * subQuantizers_ * utils::pow2(bitsPerCode_);
422 size_t smemPerBlock = getMaxSharedMemPerBlock(
device_);
424 FAISS_THROW_IF_NOT_FMT(requiredSmemSize
425 <= getMaxSharedMemPerBlock(
device_),
426 "Device %d has %zu bytes of shared memory, while "
427 "%d bits per code and %d sub-quantizers requires %zu "
428 "bytes. Consider useFloat16LookupTables and/or "
430 device_, smemPerBlock, bitsPerCode_, subQuantizers_,
437 this->d / subQuantizers_),
438 "Number of dimensions per sub-quantizer (%d) "
439 "is not currently supported without precomputed codes. "
440 "Only 1, 2, 3, 4, 6, 8, 10, 12, 16, 20, 24, 28, 32 dims "
441 "per sub-quantizer are currently supported with no "
442 "precomputed codes. "
443 "Precomputed codes supports any number of dimensions, but "
444 "will involve memory overheads.",
445 this->d / subQuantizers_);
448 FAISS_THROW_IF_NOT_MSG(this->
metric_type == faiss::METRIC_L2,
449 "METRIC_INNER_PRODUCT is currently unsupported");
std::vector< long > getListIndices(int listId) const
void precompute_table()
build precomputed table
size_t nbits
number of bits per quantization index
bool usePrecomputedTables
simple (default) implementation as an array of inverted lists
PolysemousTraining * polysemous_training
if NULL, use default
void addImpl_(int n, const float *x, const Index::idx_t *ids) override
Called from GpuIndex for add/add_with_ids.
virtual const idx_t * get_ids(size_t list_no) const =0
GpuIndexIVFPQ(GpuResources *resources, const faiss::IndexIVFPQ *index, GpuIndexIVFPQConfig config=GpuIndexIVFPQConfig())
int getListLength(int listId) const
FlatIndex * getGpuData()
For internal access.
void assign(idx_t n, const float *x, idx_t *labels, idx_t k=1)
virtual size_t list_size(size_t list_no) const =0
get the size of a list
void reserveMemory(size_t numVecs)
Reserve GPU memory in our inverted lists for this number of vectors.
int getListLength(int listId) const
bool do_polysemous_training
reorder PQ centroids after training?
size_t scan_table_threshold
use table computation or on-the-fly?
std::vector< float > precomputed_table
int polysemous_ht
Hamming thresh for polysemous filtering.
int getBitsPerCode() const
Return the number of bits per PQ code.
bool useFloat16LookupTables
virtual cudaStream_t getDefaultStream(int device)=0
long idx_t
all indices are this type
void train(Index::idx_t n, const float *x) override
static bool isSupportedPQCodeLength(int size)
Returns true if we support PQ in this size.
int nprobe_
Number of inverted list probes per query.
void reserveMemory(size_t numVecs)
Reserve GPU memory in our inverted lists for this number of vectors.
int classifyAndAddVectors(Tensor< float, 2, true > &vecs, Tensor< long, 1, true > &indices)
void query(Tensor< float, 2, true > &queries, int nprobe, int k, Tensor< float, 2, true > &outDistances, Tensor< long, 2, true > &outIndices)
const int device_
The GPU device we are resident on.
void copyFrom(const faiss::IndexIVFPQ *index)
Tensor< float, 3, true > getPQCentroids()
GpuResources * resources_
Manages streams, cuBLAS handles and scratch memory for devices.
void copyTo(faiss::IndexIVF *index) const
Copy what we have to the CPU equivalent.
void replace_invlists(InvertedLists *il, bool own=false)
replace the inverted lists, old one is deallocated if own_invlists
int nlist_
Number of inverted lists that we manage.
idx_t ntotal
total nb of indexed vectors
void searchImpl_(int n, const float *x, int k, float *distances, Index::idx_t *labels) const override
Called from GpuIndex for search.
bool verbose
verbosity level
void setPrecomputedCodes(bool enable)
Enable or disable pre-computed codes.
std::vector< unsigned char > getListCodes(int listId) const
Return the list codes of a particular list back to the CPU.
void copyTo(faiss::IndexIVFPQ *index) const
size_t nlist
number of possible key values
const MemorySpace memorySpace_
The memory space of our primary storage on the GPU.
bool by_residual
Encode residual or plain vector?
GpuIndexFlat * quantizer_
Quantizer for inverted lists.
MetricType metric_type
type of metric this index uses for search
ProductQuantizer pq
produces the codes
InvertedLists * invlists
Acess to the actual data.
size_t M
number of subquantizers
int getNumSubQuantizers() const
Return the number of sub-quantizers we are using.
std::vector< long > getListIndices(int listId) const
Return the list indices of a particular list back to the CPU.
int getCentroidsPerSubQuantizer() const
Return the number of centroids per PQ code (2^bits per code)
void setPrecomputedCodes(bool enable)
Enable or disable pre-computed codes.
virtual const uint8_t * get_codes(size_t list_no) const =0
void copyFrom(const faiss::IndexIVF *index)
Copy what we need from the CPU equivalent.
bool is_trained
set if the Index does not require training, or if training is done already
void compute_residual(const float *x, float *residual, idx_t key) const
size_t max_codes
max nb of codes to visit to do a query
bool getPrecomputedCodes() const
Are pre-computed codes enabled?
IndicesOptions indicesOptions
Index storage options for the GPU.
Implementing class for IVFPQ on the GPU.
size_t code_size
code size per vector in bytes
MetricType
Some algorithms support both an inner product version and a L2 search version.
int use_precomputed_table
if by_residual, build precompute tables
std::vector< unsigned char > getListCodes(int listId) const
std::vector< float > centroids
Centroid table, size M * ksub * dsub.
static bool isSupportedNoPrecomputedSubDimSize(int dims)