All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
GpuIndexIVFPQ.h
/**
* 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.
#pragma once
#include "GpuIndexIVF.h"
#include <vector>
namespace faiss { struct IndexIVFPQ; }
namespace faiss { namespace gpu {
class GpuIndexFlat;
class IVFPQ;
struct GpuIndexIVFPQConfig : public GpuIndexIVFConfig {
inline GpuIndexIVFPQConfig()
: useFloat16LookupTables(false),
usePrecomputedTables(false) {
}
/// Whether or not float16 residual distance tables are used in the
/// list scanning kernels. When subQuantizers * 2^bitsPerCode >
/// 16384, this is required.
bool useFloat16LookupTables;
/// Whether or not we enable the precomputed table option for
/// search, which can substantially increase the memory requirement.
bool usePrecomputedTables;
};
/// IVFPQ index for the GPU
class GpuIndexIVFPQ : public GpuIndexIVF {
public:
/// Construct from a pre-existing faiss::IndexIVFPQ instance, copying
/// data over to the given GPU, if the input index is trained.
GpuIndexIVFPQ(GpuResources* resources,
const faiss::IndexIVFPQ* index,
GpuIndexIVFPQConfig config = GpuIndexIVFPQConfig());
/// Construct an empty index
GpuIndexIVFPQ(GpuResources* resources,
int dims,
int nlist,
int subQuantizers,
int bitsPerCode,
faiss::MetricType metric,
GpuIndexIVFPQConfig config = GpuIndexIVFPQConfig());
~GpuIndexIVFPQ() override;
/// Reserve space on the GPU for the inverted lists for `num`
/// vectors, assumed equally distributed among
/// Initialize ourselves from the given CPU index; will overwrite
/// all data in ourselves
void copyFrom(const faiss::IndexIVFPQ* index);
/// Copy ourselves to the given CPU index; will overwrite all data
/// in the index instance
void copyTo(faiss::IndexIVFPQ* index) const;
/// Reserve GPU memory in our inverted lists for this number of vectors
void reserveMemory(size_t numVecs);
/// Enable or disable pre-computed codes
void setPrecomputedCodes(bool enable);
/// Are pre-computed codes enabled?
bool getPrecomputedCodes() const;
/// Return the number of sub-quantizers we are using
int getNumSubQuantizers() const;
/// Return the number of bits per PQ code
int getBitsPerCode() const;
/// Return the number of centroids per PQ code (2^bits per code)
int getCentroidsPerSubQuantizer() const;
/// After adding vectors, one can call this to reclaim device memory
/// to exactly the amount needed. Returns space reclaimed in bytes
size_t reclaimMemory();
/// Clears out all inverted lists, but retains the coarse and
/// product centroid information
void reset() override;
void train(Index::idx_t n, const float* x) override;
/// For debugging purposes, return the list length of a particular
/// list
int getListLength(int listId) const;
/// For debugging purposes, return the list codes of a particular
/// list
std::vector<unsigned char> getListCodes(int listId) const;
/// For debugging purposes, return the list indices of a particular
/// list
std::vector<long> getListIndices(int listId) const;
protected:
/// Called from GpuIndex for add/add_with_ids
void addImpl_(
faiss::Index::idx_t n,
const float* x,
const faiss::Index::idx_t* ids) override;
/// Called from GpuIndex for search
void searchImpl_(
faiss::Index::idx_t n,
const float* x,
faiss::Index::idx_t k,
float* distances,
faiss::Index::idx_t* labels) const override;
private:
void verifySettings_() const;
void trainResidualQuantizer_(Index::idx_t n, const float* x);
private:
GpuIndexIVFPQConfig ivfpqConfig_;
/// Number of sub-quantizers per encoded vector
int subQuantizers_;
/// Bits per sub-quantizer code
int bitsPerCode_;
/// Desired inverted list memory reservation
size_t reserveMemoryVecs_;
/// The product quantizer instance that we own; contains the
/// inverted lists
IVFPQ* index_;
};
} } // namespace
faiss::gpu::GpuIndexIVFPQ::getListIndices
std::vector< long > getListIndices(int listId) const
Definition: GpuIndexIVFPQ.cu:408
faiss::gpu::GpuIndexIVFPQ::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: GpuIndexIVFPQ.cu:349
faiss::gpu::GpuIndexIVFPQ::reset
void reset() override
Definition: GpuIndexIVFPQ.cu:245
faiss::gpu::GpuIndexIVFPQConfig::usePrecomputedTables
bool usePrecomputedTables
Definition: GpuIndexIVFPQ.h:36
faiss::gpu::GpuIndexIVFPQ::GpuIndexIVFPQ
GpuIndexIVFPQ(GpuResources *resources, const faiss::IndexIVFPQ *index, GpuIndexIVFPQConfig config=GpuIndexIVFPQConfig())
Definition: GpuIndexIVFPQ.cu:25
faiss::gpu::GpuIndexIVFPQ::getListLength
int getListLength(int listId) const
Definition: GpuIndexIVFPQ.cu:394
faiss::gpu::GpuIndexIVFConfig
Definition: GpuIndexIVF.h:25
faiss::gpu::GpuResources
Definition: GpuResources.h:23
faiss::gpu::GpuIndexIVFPQ::getBitsPerCode
int getBitsPerCode() const
Return the number of bits per PQ code.
Definition: GpuIndexIVFPQ.cu:225
faiss::gpu::GpuIndexIVFPQConfig::useFloat16LookupTables
bool useFloat16LookupTables
Definition: GpuIndexIVFPQ.h:32
faiss::gpu::GpuIndexIVF
Definition: GpuIndexIVF.h:37
faiss::gpu::GpuIndexIVFPQ::train
void train(Index::idx_t n, const float *x) override
Definition: GpuIndexIVFPQ.cu:301
faiss::gpu::GpuIndexIVFPQ::reclaimMemory
size_t reclaimMemory()
Definition: GpuIndexIVFPQ.cu:235
faiss::gpu::GpuIndexIVFPQConfig
Definition: GpuIndexIVFPQ.h:23
faiss::IndexIVFPQ
Definition: IndexIVFPQ.h:29
faiss::gpu::GpuIndexIVFPQ::reserveMemory
void reserveMemory(size_t numVecs)
Reserve GPU memory in our inverted lists for this number of vectors.
Definition: GpuIndexIVFPQ.cu:195
faiss::gpu::GpuIndexIVFPQ::copyFrom
void copyFrom(const faiss::IndexIVFPQ *index)
Definition: GpuIndexIVFPQ.cu:80
faiss::Index::idx_t
long idx_t
all indices are this type
Definition: Index.h:62
faiss::gpu::GpuIndexIVFPQ::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: GpuIndexIVFPQ.cu:320
faiss::gpu::GpuIndexIVFPQ::copyTo
void copyTo(faiss::IndexIVFPQ *index) const
Definition: GpuIndexIVFPQ.cu:145
faiss::gpu::GpuIndexIVFPQ
IVFPQ index for the GPU.
Definition: GpuIndexIVFPQ.h:40
faiss::gpu::GpuIndexIVFPQ::getNumSubQuantizers
int getNumSubQuantizers() const
Return the number of sub-quantizers we are using.
Definition: GpuIndexIVFPQ.cu:220
faiss::gpu::GpuIndexIVFPQ::getCentroidsPerSubQuantizer
int getCentroidsPerSubQuantizer() const
Return the number of centroids per PQ code (2^bits per code)
Definition: GpuIndexIVFPQ.cu:230
faiss::gpu::GpuIndexIVFPQ::setPrecomputedCodes
void setPrecomputedCodes(bool enable)
Enable or disable pre-computed codes.
Definition: GpuIndexIVFPQ.cu:204
faiss::gpu::GpuIndexIVFPQ::getPrecomputedCodes
bool getPrecomputedCodes() const
Are pre-computed codes enabled?
Definition: GpuIndexIVFPQ.cu:215
faiss::gpu::IVFPQ
Implementing class for IVFPQ on the GPU.
Definition: IVFPQ.cuh:19
faiss::MetricType
MetricType
Some algorithms support both an inner product vetsion and a L2 search version.
Definition: Index.h:43
faiss::gpu::GpuIndexIVFPQ::getListCodes
std::vector< unsigned char > getListCodes(int listId) const
Definition: GpuIndexIVFPQ.cu:400