Namespace faiss::gpu

namespace gpu

Enums

enum class DistanceDataType

Values:

enumerator F32

enumerator F16

enum class IndicesDataType

Values:

enumerator I64

enumerator I32

enum IndicesOptions

How user vector index data is stored on the GPU.

Values:

enumerator INDICES_CPU

The user indices are only stored on the CPU; the GPU returns (inverted list, offset) to the CPU which is then translated to the real user index.

enumerator INDICES_IVF

The indices are not stored at all, on either the CPU or GPU. Only (inverted list, offset) is returned to the user as the index.

enumerator INDICES_32_BIT

Indices are stored as 32 bit integers on the GPU, but returned as 64 bit integers

enumerator INDICES_64_BIT

Indices are stored as 64 bit integers on the GPU.

enum AllocType

Values:

enumerator Other

Unknown allocation type or miscellaneous (not currently categorized)

enumerator FlatData

Primary data storage for GpuIndexFlat (the raw matrix of vectors and vector norms if needed)

enumerator IVFLists

Primary data storage for GpuIndexIVF* (the storage for each individual IVF list)

enumerator Quantizer

Quantizer (PQ, SQ) dictionary information.

enumerator QuantizerPrecomputedCodes

For GpuIndexIVFPQ, “precomputed codes” for more efficient PQ lookup require the use of possibly large tables. These are marked separately from Quantizer as these can frequently be 100s - 1000s of MiB in size

enumerator TemporaryMemoryBuffer

StandardGpuResources implementation specific types When using StandardGpuResources, temporary memory allocations (MemorySpace::Temporary) come out of a stack region of memory that is allocated up front for each gpu (e.g., 1.5 GiB upon initialization). This allocation by StandardGpuResources is marked with this AllocType.

enumerator TemporaryMemoryOverflow

When using StandardGpuResources, any MemorySpace::Temporary allocations that cannot be satisfied within the TemporaryMemoryBuffer region fall back to calling cudaMalloc which are sized to just the request at hand. These “overflow” temporary allocations are marked with this AllocType.

enum MemorySpace

Memory regions accessible to the GPU.

Values:

enumerator Temporary

Temporary device memory (guaranteed to no longer be used upon exit of a top-level index call, and where the streams using it have completed GPU work). Typically backed by Device memory (cudaMalloc/cudaFree).

enumerator Device

Managed using cudaMalloc/cudaFree (typical GPU device memory)

enumerator Unified

Managed using cudaMallocManaged/cudaFree (typical Unified CPU/GPU memory)

Functions

faiss::Index *index_gpu_to_cpu(const faiss::Index *gpu_index)

converts any GPU index inside gpu_index to a CPU index

faiss::Index *index_cpu_to_gpu(GpuResourcesProvider *provider, int device, const faiss::Index *index, const GpuClonerOptions *options = nullptr)

converts any CPU index that can be converted to GPU

faiss::Index *index_cpu_to_gpu_multiple(std::vector<GpuResourcesProvider*> &provider, std::vector<int> &devices, const faiss::Index *index, const GpuMultipleClonerOptions *options = nullptr)

void bfKnn(GpuResourcesProvider *resources, const GpuDistanceParams &args)

A wrapper for gpu/impl/Distance.cuh to expose direct brute-force k-nearest neighbor searches on an externally-provided region of memory (e.g., from a pytorch tensor). The data (vectors, queries, outDistances, outIndices) can be resident on the GPU or the CPU, but all calculations are performed on the GPU. If the result buffers are on the CPU, results will be copied back when done.

All GPU computation is performed on the current CUDA device, and ordered with respect to resources->getDefaultStreamCurrentDevice().

For each vector in queries, searches all of vectors to find its k nearest neighbors with respect to the given metric

void bruteForceKnn(GpuResourcesProvider *resources, faiss::MetricType metric, const float *vectors, bool vectorsRowMajor, int numVectors, const float *queries, bool queriesRowMajor, int numQueries, int dims, int k, float *outDistances, idx_t *outIndices)

Deprecated legacy implementation.

GpuIndex *tryCastGpuIndex(faiss::Index *index)

If the given index is a GPU index, this returns the index instance.

bool isGpuIndex(faiss::Index *index)

Is the given index instance a GPU index?

bool isGpuIndexImplemented(faiss::Index *index)

Does the given CPU index instance have a corresponding GPU implementation?

std::string allocTypeToString(AllocType t)

Convert an AllocType to string.

std::string memorySpaceToString(MemorySpace s)

Convert a MemorySpace to string.

AllocInfo makeDevAlloc(AllocType at, cudaStream_t st)

Create an AllocInfo for the current device with MemorySpace::Device.

AllocInfo makeTempAlloc(AllocType at, cudaStream_t st)

Create an AllocInfo for the current device with MemorySpace::Temporary.

AllocInfo makeSpaceAlloc(AllocType at, MemorySpace sp, cudaStream_t st)

Create an AllocInfo for the current device.

void newTestSeed()

Generates and displays a new seed for the test.

void setTestSeed(long seed)

Uses an explicit seed for the test.

float relativeError(float a, float b)

Returns the relative error in difference between a and b (|a - b| / (0.5 * (|a| + |b|))

int randVal(int a, int b)

Generates a random integer in the range [a, b].

bool randBool()

Generates a random bool.

template<typename T>
T randSelect(std::initializer_list<T> vals)

Select a random value from the given list of values provided as an initializer_list

std::vector<float> randVecs(size_t num, size_t dim)

Generates a collection of random vectors in the range [0, 1].

std::vector<unsigned char> randBinaryVecs(size_t num, size_t dim)

Generates a collection of random bit vectors.

std::vector<float> roundToHalf(const std::vector<float> &v)

void compareIndices(const std::vector<float> &queryVecs, faiss::Index &refIndex, faiss::Index &testIndex, int numQuery, int dim, int k, const std::string &configMsg, float maxRelativeError = 6e-5f, float pctMaxDiff1 = 0.1f, float pctMaxDiffN = 0.005f)

Compare two indices via query for similarity, with a user-specified set of query vectors

void compareIndices(faiss::Index &refIndex, faiss::Index &testIndex, int numQuery, int dim, int k, const std::string &configMsg, float maxRelativeError = 6e-5f, float pctMaxDiff1 = 0.1f, float pctMaxDiffN = 0.005f)

Compare two indices via query for similarity, generating random query vectors

void compareLists(const float *refDist, const faiss::idx_t *refInd, const float *testDist, const faiss::idx_t *testInd, int dim1, int dim2, const std::string &configMsg, bool printBasicStats, bool printDiffs, bool assertOnErr, float maxRelativeError = 6e-5f, float pctMaxDiff1 = 0.1f, float pctMaxDiffN = 0.005f)

Display specific differences in the two (distance, index) lists.

template<typename A, typename B>
void testIVFEquality(A &cpuIndex, B &gpuIndex)

Compare IVF lists between a CPU and GPU index.

int getCurrentDevice()

Returns the current thread-local GPU device.

void setCurrentDevice(int device)

Sets the current thread-local GPU device.

int getNumDevices()

Returns the number of available GPU devices.

void profilerStart()

Starts the CUDA profiler (exposed via SWIG)

void profilerStop()

Stops the CUDA profiler (exposed via SWIG)

void synchronizeAllDevices()

Synchronizes the CPU against all devices (equivalent to cudaDeviceSynchronize for each device)

const cudaDeviceProp &getDeviceProperties(int device)

Returns a cached cudaDeviceProp for the given device.

const cudaDeviceProp &getCurrentDeviceProperties()

Returns the cached cudaDeviceProp for the current device.

int getMaxThreads(int device)

Returns the maximum number of threads available for the given GPU device

int getMaxThreadsCurrentDevice()

Equivalent to getMaxThreads(getCurrentDevice())

dim3 getMaxGrid(int device)

Returns the maximum grid size for the given GPU device.

dim3 getMaxGridCurrentDevice()

Equivalent to getMaxGrid(getCurrentDevice())

size_t getMaxSharedMemPerBlock(int device)

Returns the maximum smem available for the given GPU device.

size_t getMaxSharedMemPerBlockCurrentDevice()

Equivalent to getMaxSharedMemPerBlock(getCurrentDevice())

int getDeviceForAddress(const void *p)

For a given pointer, returns whether or not it is located on a device (deviceId >= 0) or the host (-1).

bool getFullUnifiedMemSupport(int device)

Does the given device support full unified memory sharing host memory?

bool getFullUnifiedMemSupportCurrentDevice()

Equivalent to getFullUnifiedMemSupport(getCurrentDevice())

bool getTensorCoreSupport(int device)

Does the given device support tensor core operations?

bool getTensorCoreSupportCurrentDevice()

Equivalent to getTensorCoreSupport(getCurrentDevice())

size_t getFreeMemory(int device)

Returns the amount of currently available memory on the given device.

size_t getFreeMemoryCurrentDevice()

Equivalent to getFreeMemory(getCurrentDevice())

template<typename L1, typename L2>
void streamWaitBase(const L1 &listWaiting, const L2 &listWaitOn)

Call for a collection of streams to wait on.

template<typename L1>
void streamWait(const L1 &a, const std::initializer_list<cudaStream_t> &b)

These versions allow usage of initializer_list as arguments, since otherwise {…} doesn’t have a type

template<typename L2>
void streamWait(const std::initializer_list<cudaStream_t> &a, const L2 &b)

inline void streamWait(const std::initializer_list<cudaStream_t> &a, const std::initializer_list<cudaStream_t> &b)

struct GpuParameterSpace : public faiss::ParameterSpace

#include <GpuAutoTune.h>

parameter space and setters for GPU indexes

Public Functions

virtual void initialize(const faiss::Index *index) override

initialize with reasonable parameters for the index

virtual void set_index_parameter(faiss::Index *index, const std::string &name, double val) const override

set a combination of parameters on an index

size_t n_combinations() const

nb of combinations, = product of values sizes

bool combination_ge(size_t c1, size_t c2) const

returns whether combinations c1 >= c2 in the tuple sense

std::string combination_name(size_t cno) const

get string representation of the combination

void display() const

print a description on stdout

ParameterRange &add_range(const std::string &name)

add a new parameter (or return it if it exists)

void set_index_parameters(Index *index, size_t cno) const

set a combination of parameters on an index

void set_index_parameters(Index *index, const char *param_string) const

set a combination of parameters described by a string

void update_bounds(size_t cno, const OperatingPoint &op, double *upper_bound_perf, double *lower_bound_t) const

find an upper bound on the performance and a lower bound on t for configuration cno given another operating point op

void explore(Index *index, size_t nq, const float *xq, const AutoTuneCriterion &crit, OperatingPoints *ops) const

explore operating points

Parameters:

• index – index to run on

• xq – query vectors (size nq * index.d)

• crit – selection criterion

• ops – resulting operating points

Public Members

std::vector<ParameterRange> parameter_ranges

all tunable parameters

int verbose

verbosity during exploration

int n_experiments

nb of experiments during optimization (0 = try all combinations)

size_t batchsize

maximum number of queries to submit at a time.

bool thread_over_batches

use multithreading over batches (useful to benchmark independent single-searches)

double min_test_duration

run tests several times until they reach at least this duration (to avoid jittering in MT mode)

struct ToCPUCloner : public faiss::Cloner

#include <GpuCloner.h>

Cloner specialized for GPU -> CPU.

Public Functions

void merge_index(Index *dst, Index *src, bool successive_ids)

virtual Index *clone_Index(const Index *index) override

virtual VectorTransform *clone_VectorTransform(const VectorTransform*)

virtual IndexIVF *clone_IndexIVF(const IndexIVF*)

struct ToGpuCloner : public faiss::Cloner, public faiss::gpu::GpuClonerOptions

#include <GpuCloner.h>

Cloner specialized for CPU -> 1 GPU.

Public Functions

ToGpuCloner(GpuResourcesProvider *prov, int device, const GpuClonerOptions &options)

virtual Index *clone_Index(const Index *index) override

virtual VectorTransform *clone_VectorTransform(const VectorTransform*)

virtual IndexIVF *clone_IndexIVF(const IndexIVF*)

Public Members

GpuResourcesProvider *provider

int device

IndicesOptions indicesOptions = INDICES_64_BIT

how should indices be stored on index types that support indices (anything but GpuIndexFlat*)?

bool useFloat16CoarseQuantizer = false

is the coarse quantizer in float16?

bool useFloat16 = false

for GpuIndexIVFFlat, is storage in float16? for GpuIndexIVFPQ, are intermediate calculations in float16?

bool usePrecomputed = false

use precomputed tables?

long reserveVecs = 0

reserve vectors in the invfiles?

bool storeTransposed = false

For GpuIndexFlat, store data in transposed layout?

bool verbose = false

Set verbose options on the index.

struct ToGpuClonerMultiple : public faiss::Cloner, public faiss::gpu::GpuMultipleClonerOptions

#include <GpuCloner.h>

Cloner specialized for CPU -> multiple GPUs.

Public Functions

ToGpuClonerMultiple(std::vector<GpuResourcesProvider*> &provider, std::vector<int> &devices, const GpuMultipleClonerOptions &options)

ToGpuClonerMultiple(const std::vector<ToGpuCloner> &sub_cloners, const GpuMultipleClonerOptions &options)

void copy_ivf_shard(const IndexIVF *index_ivf, IndexIVF *idx2, long n, long i)

Index *clone_Index_to_shards(const Index *index)

virtual Index *clone_Index(const Index *index) override

main function

virtual VectorTransform *clone_VectorTransform(const VectorTransform*)

virtual IndexIVF *clone_IndexIVF(const IndexIVF*)

Public Members

std::vector<ToGpuCloner> sub_cloners

bool shard = false

Whether to shard the index across GPUs, versus replication across GPUs

int shard_type = 1

IndexIVF::copy_subset_to subset type.

IndicesOptions indicesOptions = INDICES_64_BIT

how should indices be stored on index types that support indices (anything but GpuIndexFlat*)?

bool useFloat16CoarseQuantizer = false

is the coarse quantizer in float16?

bool useFloat16 = false

for GpuIndexIVFFlat, is storage in float16? for GpuIndexIVFPQ, are intermediate calculations in float16?

bool usePrecomputed = false

use precomputed tables?

long reserveVecs = 0

reserve vectors in the invfiles?

bool storeTransposed = false

For GpuIndexFlat, store data in transposed layout?

bool verbose = false

Set verbose options on the index.

struct GpuProgressiveDimIndexFactory : public faiss::ProgressiveDimIndexFactory

#include <GpuCloner.h>

index factory for the ProgressiveDimClustering object

Public Functions

explicit GpuProgressiveDimIndexFactory(int ngpu)

virtual Index *operator()(int dim) override

ownership transferred to caller

virtual ~GpuProgressiveDimIndexFactory() override

Public Members

GpuMultipleClonerOptions options

std::vector<GpuResourcesProvider*> vres

std::vector<int> devices

int ncall

struct GpuClonerOptions

#include <GpuClonerOptions.h>

set some options on how to copy to GPU

Subclassed by faiss::gpu::GpuMultipleClonerOptions, faiss::gpu::ToGpuCloner

Public Functions

GpuClonerOptions()

Public Members

IndicesOptions indicesOptions = INDICES_64_BIT

how should indices be stored on index types that support indices (anything but GpuIndexFlat*)?

bool useFloat16CoarseQuantizer = false

is the coarse quantizer in float16?

bool useFloat16 = false

for GpuIndexIVFFlat, is storage in float16? for GpuIndexIVFPQ, are intermediate calculations in float16?

bool usePrecomputed = false

use precomputed tables?

long reserveVecs = 0

reserve vectors in the invfiles?

bool storeTransposed = false

For GpuIndexFlat, store data in transposed layout?

bool verbose = false

Set verbose options on the index.

struct GpuMultipleClonerOptions : public faiss::gpu::GpuClonerOptions

Subclassed by faiss::gpu::ToGpuClonerMultiple

Public Functions

GpuMultipleClonerOptions()

Public Members

bool shard = false

Whether to shard the index across GPUs, versus replication across GPUs

int shard_type = 1

IndexIVF::copy_subset_to subset type.

IndicesOptions indicesOptions = INDICES_64_BIT

how should indices be stored on index types that support indices (anything but GpuIndexFlat*)?

bool useFloat16CoarseQuantizer = false

is the coarse quantizer in float16?

bool useFloat16 = false

for GpuIndexIVFFlat, is storage in float16? for GpuIndexIVFPQ, are intermediate calculations in float16?

bool usePrecomputed = false

use precomputed tables?

long reserveVecs = 0

reserve vectors in the invfiles?

bool storeTransposed = false

For GpuIndexFlat, store data in transposed layout?

bool verbose = false

Set verbose options on the index.

struct GpuDistanceParams

#include <GpuDistance.h>

Arguments to brute-force GPU k-nearest neighbor searching.

Public Functions

inline GpuDistanceParams()

Public Members

faiss::MetricType metric

Search parameter: distance metric.

float metricArg

Search parameter: distance metric argument (if applicable) For metric == METRIC_Lp, this is the p-value

int k

Search parameter: return k nearest neighbors If the value provided is -1, then we report all pairwise distances without top-k filtering

int dims

Vector dimensionality.

const void *vectors

If vectorsRowMajor is true, this is numVectors x dims, with dims innermost; otherwise, dims x numVectors, with numVectors innermost

DistanceDataType vectorType

bool vectorsRowMajor

int numVectors

const float *vectorNorms

Precomputed L2 norms for each vector in vectors, which can be optionally provided in advance to speed computation for METRIC_L2

const void *queries

If queriesRowMajor is true, this is numQueries x dims, with dims innermost; otherwise, dims x numQueries, with numQueries innermost

DistanceDataType queryType

bool queriesRowMajor

int numQueries

float *outDistances

A region of memory size numQueries x k, with k innermost (row major) if k > 0, or if k == -1, a region of memory of size numQueries x numVectors

bool ignoreOutDistances

Do we only care about the indices reported, rather than the output distances? Not used if k == -1 (all pairwise distances)

IndicesDataType outIndicesType

A region of memory size numQueries x k, with k innermost (row major). Not used if k == -1 (all pairwise distances)

void *outIndices

int device

On which GPU device should the search run? -1 indicates that the current CUDA thread-local device (via cudaGetDevice/cudaSetDevice) is used Otherwise, an integer 0 <= device < numDevices indicates the device for execution

class GpuIcmEncoder : public lsq::IcmEncoder

#include <GpuIcmEncoder.h>

Perform LSQ encoding on GPU.

Split input vectors to different devices and call IcmEncoderImpl::encode to encode them

Public Functions

GpuIcmEncoder(const LocalSearchQuantizer *lsq, const std::vector<GpuResourcesProvider*> &provs, const std::vector<int> &devices)

~GpuIcmEncoder()

GpuIcmEncoder(const GpuIcmEncoder&) = delete

GpuIcmEncoder &operator=(const GpuIcmEncoder&) = delete

void set_binary_term() override

void encode(int32_t *codes, const float *x, std::mt19937 &gen, size_t n, size_t ils_iters) const override

Private Members

std::unique_ptr<IcmEncoderShards> shards

struct GpuIcmEncoderFactory : public lsq::IcmEncoderFactory

Public Functions

explicit GpuIcmEncoderFactory(int ngpus = 1)

lsq::IcmEncoder *get(const LocalSearchQuantizer *lsq) override

Public Members

std::vector<GpuResourcesProvider*> provs

std::vector<int> devices

struct GpuIndexConfig

Subclassed by faiss::gpu::GpuIndexBinaryFlatConfig, faiss::gpu::GpuIndexFlatConfig, faiss::gpu::GpuIndexIVFConfig

Public Functions

inline GpuIndexConfig()

Public Members

int device

GPU device on which the index is resident.

MemorySpace memorySpace

What memory space to use for primary storage. On Pascal and above (CC 6+) architectures, allows GPUs to use more memory than is available on the GPU.

class GpuIndex : public faiss::Index

Subclassed by faiss::gpu::GpuIndexFlat, faiss::gpu::GpuIndexIVF

Public Types

using component_t = float

using distance_t = float

Public Functions

GpuIndex(std::shared_ptr<GpuResources> resources, int dims, faiss::MetricType metric, float metricArg, GpuIndexConfig config)

int getDevice() const

Returns the device that this index is resident on.

std::shared_ptr<GpuResources> getResources()

Returns a reference to our GpuResources object that manages memory, stream and handle resources on the GPU

void setMinPagingSize(size_t size)

Set the minimum data size for searches (in MiB) for which we use CPU -> GPU paging

size_t getMinPagingSize() const

Returns the current minimum data size for paged searches.

virtual void add(idx_t, const float *x) override

x can be resident on the CPU or any GPU; copies are performed as needed Handles paged adds if the add set is too large; calls addInternal_

virtual void add_with_ids(idx_t n, const float *x, const idx_t *ids) override

x and ids can be resident on the CPU or any GPU; copies are performed as needed Handles paged adds if the add set is too large; calls addInternal_

virtual void assign(idx_t n, const float *x, idx_t *labels, idx_t k = 1) const override

x and labels can be resident on the CPU or any GPU; copies are performed as needed

virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels, const SearchParameters *params = nullptr) const override

x, distances and labels can be resident on the CPU or any GPU; copies are performed as needed

virtual void search_and_reconstruct(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels, float *recons, const SearchParameters *params = nullptr) const override

x, distances and labels and recons can be resident on the CPU or any GPU; copies are performed as needed

virtual void compute_residual(const float *x, float *residual, idx_t key) const override

Overridden to force GPU indices to provide their own GPU-friendly implementation

virtual void compute_residual_n(idx_t n, const float *xs, float *residuals, const idx_t *keys) const override

Overridden to force GPU indices to provide their own GPU-friendly implementation

virtual void train(idx_t n, const float *x)

Perform training on a representative set of vectors

Parameters:

• n – nb of training vectors

• x – training vecors, size n * d

virtual void range_search(idx_t n, const float *x, float radius, RangeSearchResult *result, const SearchParameters *params = nullptr) const

query n vectors of dimension d to the index.

return all vectors with distance < radius. Note that many indexes do not implement the range_search (only the k-NN search is mandatory).

Parameters:

• x – input vectors to search, size n * d

• radius – search radius

• result – result table

virtual void reset() = 0

removes all elements from the database.

virtual size_t remove_ids(const IDSelector &sel)

removes IDs from the index. Not supported by all indexes. Returns the number of elements removed.

virtual void reconstruct(idx_t key, float *recons) const

Reconstruct a stored vector (or an approximation if lossy coding)

this function may not be defined for some indexes

Parameters:

• key – id of the vector to reconstruct

• recons – reconstucted vector (size d)

virtual void reconstruct_batch(idx_t n, const idx_t *keys, float *recons) const

Reconstruct several stored vectors (or an approximation if lossy coding)

this function may not be defined for some indexes

Parameters:

• n – number of vectors to reconstruct

• keys – ids of the vectors to reconstruct (size n)

• recons – reconstucted vector (size n * d)

virtual void reconstruct_n(idx_t i0, idx_t ni, float *recons) const

Reconstruct vectors i0 to i0 + ni - 1

this function may not be defined for some indexes

Parameters:

recons – reconstucted vector (size ni * d)

virtual DistanceComputer *get_distance_computer() const

Get a DistanceComputer (defined in AuxIndexStructures) object for this kind of index.

DistanceComputer is implemented for indexes that support random access of their vectors.

virtual size_t sa_code_size() const

size of the produced codes in bytes

virtual void sa_encode(idx_t n, const float *x, uint8_t *bytes) const

encode a set of vectors

Parameters:

• n – number of vectors

• x – input vectors, size n * d

• bytes – output encoded vectors, size n * sa_code_size()

virtual void sa_decode(idx_t n, const uint8_t *bytes, float *x) const

decode a set of vectors

Parameters:

• n – number of vectors

• bytes – input encoded vectors, size n * sa_code_size()

• x – output vectors, size n * d

virtual void merge_from(Index &otherIndex, idx_t add_id = 0)

moves the entries from another dataset to self. On output, other is empty. add_id is added to all moved ids (for sequential ids, this would be this->ntotal)

virtual void check_compatible_for_merge(const Index &otherIndex) const

check that the two indexes are compatible (ie, they are trained in the same way and have the same parameters). Otherwise throw.

Public Members

int d

vector dimension

idx_t ntotal

total nb of indexed vectors

bool verbose

verbosity level

bool is_trained

set if the Index does not require training, or if training is done already

MetricType metric_type

type of metric this index uses for search

float metric_arg

argument of the metric type

Protected Functions

void copyFrom(const faiss::Index *index)

Copy what we need from the CPU equivalent.

void copyTo(faiss::Index *index) const

Copy what we have to the CPU equivalent.

virtual bool addImplRequiresIDs_() const = 0

Does addImpl_ require IDs? If so, and no IDs are provided, we will generate them sequentially based on the order in which the IDs are added

virtual void addImpl_(int n, const float *x, const idx_t *ids) = 0

Overridden to actually perform the add All data is guaranteed to be resident on our device

virtual void searchImpl_(int n, const float *x, int k, float *distances, idx_t *labels, const SearchParameters *params) const = 0

Overridden to actually perform the search All data is guaranteed to be resident on our device

Protected Attributes

std::shared_ptr<GpuResources> resources_

Manages streams, cuBLAS handles and scratch memory for devices.

const GpuIndexConfig config_

Our configuration options.

size_t minPagedSize_

Size above which we page copies from the CPU to GPU.

Private Functions

void addPaged_(int n, const float *x, const idx_t *ids)

Handles paged adds if the add set is too large, passes to addImpl_ to actually perform the add for the current page

void addPage_(int n, const float *x, const idx_t *ids)

Calls addImpl_ for a single page of GPU-resident data.

void searchNonPaged_(int n, const float *x, int k, float *outDistancesData, idx_t *outIndicesData, const SearchParameters *params) const

Calls searchImpl_ for a single page of GPU-resident data.

void searchFromCpuPaged_(int n, const float *x, int k, float *outDistancesData, idx_t *outIndicesData, const SearchParameters *params) const

Calls searchImpl_ for a single page of GPU-resident data, handling paging of the data and copies from the CPU

struct GpuIndexBinaryFlatConfig : public faiss::gpu::GpuIndexConfig

Public Members

int device

GPU device on which the index is resident.

MemorySpace memorySpace

What memory space to use for primary storage. On Pascal and above (CC 6+) architectures, allows GPUs to use more memory than is available on the GPU.

class GpuIndexBinaryFlat : public faiss::IndexBinary

#include <GpuIndexBinaryFlat.h>

A GPU version of IndexBinaryFlat for brute-force comparison of bit vectors via Hamming distance

Public Types

using component_t = uint8_t

using distance_t = int32_t

Public Functions

GpuIndexBinaryFlat(GpuResourcesProvider *resources, const faiss::IndexBinaryFlat *index, GpuIndexBinaryFlatConfig config = GpuIndexBinaryFlatConfig())

Construct from a pre-existing faiss::IndexBinaryFlat instance, copying data over to the given GPU

GpuIndexBinaryFlat(GpuResourcesProvider *resources, int dims, GpuIndexBinaryFlatConfig config = GpuIndexBinaryFlatConfig())

Construct an empty instance that can be added to.

~GpuIndexBinaryFlat() override

int getDevice() const

Returns the device that this index is resident on.

std::shared_ptr<GpuResources> getResources()

Returns a reference to our GpuResources object that manages memory, stream and handle resources on the GPU

void copyFrom(const faiss::IndexBinaryFlat *index)

Initialize ourselves from the given CPU index; will overwrite all data in ourselves

void copyTo(faiss::IndexBinaryFlat *index) const

Copy ourselves to the given CPU index; will overwrite all data in the index instance

virtual void add(faiss::idx_t n, const uint8_t *x) override

Add n vectors of dimension d to the index.

Vectors are implicitly assigned labels ntotal .. ntotal + n - 1

Parameters:

x – input matrix, size n * d / 8

virtual void reset() override

Removes all elements from the database.

virtual void search(faiss::idx_t n, const uint8_t *x, faiss::idx_t k, int32_t *distances, faiss::idx_t *labels, const faiss::SearchParameters *params = nullptr) const override

Query n vectors of dimension d to the index.

return at most k vectors. If there are not enough results for a query, the result array is padded with -1s.

Parameters:

• x – input vectors to search, size n * d / 8

• labels – output labels of the NNs, size n*k

• distances – output pairwise distances, size n*k

virtual void reconstruct(faiss::idx_t key, uint8_t *recons) const override

Reconstruct a stored vector.

This function may not be defined for some indexes.

Parameters:

• key – id of the vector to reconstruct

• recons – reconstucted vector (size d / 8)

virtual void train(idx_t n, const uint8_t *x)

Perform training on a representative set of vectors.

Parameters:

• n – nb of training vectors

• x – training vecors, size n * d / 8

virtual void add_with_ids(idx_t n, const uint8_t *x, const idx_t *xids)

Same as add, but stores xids instead of sequential ids.

The default implementation fails with an assertion, as it is not supported by all indexes.

Parameters:

xids – if non-null, ids to store for the vectors (size n)

virtual void range_search(idx_t n, const uint8_t *x, int radius, RangeSearchResult *result, const SearchParameters *params = nullptr) const

Query n vectors of dimension d to the index.

return all vectors with distance < radius. Note that many indexes do not implement the range_search (only the k-NN search is mandatory). The distances are converted to float to reuse the RangeSearchResult structure, but they are integer. By convention, only distances < radius (strict comparison) are returned, ie. radius = 0 does not return any result and 1 returns only exact same vectors.

Parameters:

• x – input vectors to search, size n * d / 8

• radius – search radius

• result – result table

void assign(idx_t n, const uint8_t *x, idx_t *labels, idx_t k = 1) const

Return the indexes of the k vectors closest to the query x.

This function is identical to search but only returns labels of neighbors.

Parameters:

• x – input vectors to search, size n * d / 8

• labels – output labels of the NNs, size n*k

virtual size_t remove_ids(const IDSelector &sel)

Removes IDs from the index. Not supported by all indexes.

virtual void reconstruct_n(idx_t i0, idx_t ni, uint8_t *recons) const

Reconstruct vectors i0 to i0 + ni - 1.

This function may not be defined for some indexes.

Parameters:

recons – reconstucted vectors (size ni * d / 8)

virtual void search_and_reconstruct(idx_t n, const uint8_t *x, idx_t k, int32_t *distances, idx_t *labels, uint8_t *recons, const SearchParameters *params = nullptr) const

Similar to search, but also reconstructs the stored vectors (or an approximation in the case of lossy coding) for the search results.

If there are not enough results for a query, the resulting array is padded with -1s.

Parameters:

recons – reconstructed vectors size (n, k, d)

void display() const

Display the actual class name and some more info.

virtual void merge_from(IndexBinary &otherIndex, idx_t add_id = 0)

moves the entries from another dataset to self. On output, other is empty. add_id is added to all moved ids (for sequential ids, this would be this->ntotal)

virtual void check_compatible_for_merge(const IndexBinary &otherIndex) const

check that the two indexes are compatible (ie, they are trained in the same way and have the same parameters). Otherwise throw.

Public Members

int d

vector dimension

int code_size

number of bytes per vector ( = d / 8 )

idx_t ntotal

total nb of indexed vectors

bool verbose

verbosity level

bool is_trained

set if the Index does not require training, or if training is done already

MetricType metric_type

type of metric this index uses for search

Protected Functions

void searchFromCpuPaged_(int n, const uint8_t *x, int k, int32_t *outDistancesData, int *outIndicesData) const

Called from search when the input data is on the CPU; potentially allows for pinned memory usage

void searchNonPaged_(int n, const uint8_t *x, int k, int32_t *outDistancesData, int *outIndicesData) const

Protected Attributes

std::shared_ptr<GpuResources> resources_

Manages streans, cuBLAS handles and scratch memory for devices.

const GpuIndexBinaryFlatConfig binaryFlatConfig_

Configuration options.

std::unique_ptr<BinaryFlatIndex> data_

Holds our GPU data containing the list of vectors.

struct GpuIndexFlatConfig : public faiss::gpu::GpuIndexConfig

Public Functions

inline GpuIndexFlatConfig()

Public Members

bool useFloat16

Whether or not data is stored as float16.

bool storeTransposed

Deprecated: no longer used Previously used to indicate whether internal storage of vectors is transposed

int device

GPU device on which the index is resident.

MemorySpace memorySpace

What memory space to use for primary storage. On Pascal and above (CC 6+) architectures, allows GPUs to use more memory than is available on the GPU.

class GpuIndexFlat : public faiss::gpu::GpuIndex

#include <GpuIndexFlat.h>

Wrapper around the GPU implementation that looks like faiss::IndexFlat; copies over centroid data from a given faiss::IndexFlat

Subclassed by faiss::gpu::GpuIndexFlatIP, faiss::gpu::GpuIndexFlatL2

Public Types

using component_t = float

using distance_t = float

Public Functions

GpuIndexFlat(GpuResourcesProvider *provider, const faiss::IndexFlat *index, GpuIndexFlatConfig config = GpuIndexFlatConfig())

Construct from a pre-existing faiss::IndexFlat instance, copying data over to the given GPU

GpuIndexFlat(std::shared_ptr<GpuResources> resources, const faiss::IndexFlat *index, GpuIndexFlatConfig config = GpuIndexFlatConfig())

GpuIndexFlat(GpuResourcesProvider *provider, int dims, faiss::MetricType metric, GpuIndexFlatConfig config = GpuIndexFlatConfig())

Construct an empty instance that can be added to.

GpuIndexFlat(std::shared_ptr<GpuResources> resources, int dims, faiss::MetricType metric, GpuIndexFlatConfig config = GpuIndexFlatConfig())

~GpuIndexFlat() override

void copyFrom(const faiss::IndexFlat *index)

Initialize ourselves from the given CPU index; will overwrite all data in ourselves

void copyTo(faiss::IndexFlat *index) const

Copy ourselves to the given CPU index; will overwrite all data in the index instance

size_t getNumVecs() const

Returns the number of vectors we contain.

virtual void reset() override

Clears all vectors from this index.

virtual void train(idx_t n, const float *x) override

This index is not trained, so this does nothing.

virtual void add(idx_t, const float *x) override

Overrides to avoid excessive copies.

virtual void reconstruct(idx_t key, float *out) const override

Reconstruction methods; prefer the batch reconstruct as it will be more efficient

virtual void reconstruct_n(idx_t i0, idx_t num, float *out) const override

Batch reconstruction method.

virtual void reconstruct_batch(idx_t n, const idx_t *keys, float *out) const override

Batch reconstruction method.

virtual void compute_residual(const float *x, float *residual, idx_t key) const override

Compute residual.

virtual void compute_residual_n(idx_t n, const float *xs, float *residuals, const idx_t *keys) const override

Compute residual (batch mode)

inline FlatIndex *getGpuData()

For internal access.

int getDevice() const

Returns the device that this index is resident on.

std::shared_ptr<GpuResources> getResources()

Returns a reference to our GpuResources object that manages memory, stream and handle resources on the GPU

void setMinPagingSize(size_t size)

Set the minimum data size for searches (in MiB) for which we use CPU -> GPU paging

size_t getMinPagingSize() const

Returns the current minimum data size for paged searches.

virtual void add_with_ids(idx_t n, const float *x, const idx_t *ids) override

x and ids can be resident on the CPU or any GPU; copies are performed as needed Handles paged adds if the add set is too large; calls addInternal_

virtual void assign(idx_t n, const float *x, idx_t *labels, idx_t k = 1) const override

x and labels can be resident on the CPU or any GPU; copies are performed as needed

virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels, const SearchParameters *params = nullptr) const override

x, distances and labels can be resident on the CPU or any GPU; copies are performed as needed

virtual void search_and_reconstruct(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels, float *recons, const SearchParameters *params = nullptr) const override

x, distances and labels and recons can be resident on the CPU or any GPU; copies are performed as needed

virtual void range_search(idx_t n, const float *x, float radius, RangeSearchResult *result, const SearchParameters *params = nullptr) const

query n vectors of dimension d to the index.

return all vectors with distance < radius. Note that many indexes do not implement the range_search (only the k-NN search is mandatory).

Parameters:

• x – input vectors to search, size n * d

• radius – search radius

• result – result table

virtual size_t remove_ids(const IDSelector &sel)

removes IDs from the index. Not supported by all indexes. Returns the number of elements removed.

virtual DistanceComputer *get_distance_computer() const

Get a DistanceComputer (defined in AuxIndexStructures) object for this kind of index.

DistanceComputer is implemented for indexes that support random access of their vectors.

virtual size_t sa_code_size() const

size of the produced codes in bytes

virtual void sa_encode(idx_t n, const float *x, uint8_t *bytes) const

encode a set of vectors

Parameters:

• n – number of vectors

• x – input vectors, size n * d

• bytes – output encoded vectors, size n * sa_code_size()

virtual void sa_decode(idx_t n, const uint8_t *bytes, float *x) const

decode a set of vectors

Parameters:

• n – number of vectors

• bytes – input encoded vectors, size n * sa_code_size()

• x – output vectors, size n * d

virtual void merge_from(Index &otherIndex, idx_t add_id = 0)

moves the entries from another dataset to self. On output, other is empty. add_id is added to all moved ids (for sequential ids, this would be this->ntotal)

virtual void check_compatible_for_merge(const Index &otherIndex) const

check that the two indexes are compatible (ie, they are trained in the same way and have the same parameters). Otherwise throw.

Public Members

int d

vector dimension

idx_t ntotal

total nb of indexed vectors

bool verbose

verbosity level

bool is_trained

set if the Index does not require training, or if training is done already

MetricType metric_type

type of metric this index uses for search

float metric_arg

argument of the metric type

Protected Functions

virtual bool addImplRequiresIDs_() const override

Flat index does not require IDs as there is no storage available for them

virtual void addImpl_(int n, const float *x, const idx_t *ids) override

Called from GpuIndex for add.

virtual void searchImpl_(int n, const float *x, int k, float *distances, idx_t *labels, const SearchParameters *params) const override

Called from GpuIndex for search.

void copyFrom(const faiss::Index *index)

Copy what we need from the CPU equivalent.

void copyTo(faiss::Index *index) const

Copy what we have to the CPU equivalent.

Protected Attributes

const GpuIndexFlatConfig flatConfig_

Our configuration options.

std::unique_ptr<FlatIndex> data_

Holds our GPU data containing the list of vectors.

std::shared_ptr<GpuResources> resources_

Manages streams, cuBLAS handles and scratch memory for devices.

const GpuIndexConfig config_

Our configuration options.

size_t minPagedSize_

Size above which we page copies from the CPU to GPU.

class GpuIndexFlatL2 : public faiss::gpu::GpuIndexFlat

#include <GpuIndexFlat.h>

Wrapper around the GPU implementation that looks like faiss::IndexFlatL2; copies over centroid data from a given faiss::IndexFlat

Public Types

using component_t = float

using distance_t = float

Public Functions

GpuIndexFlatL2(GpuResourcesProvider *provider, faiss::IndexFlatL2 *index, GpuIndexFlatConfig config = GpuIndexFlatConfig())

Construct from a pre-existing faiss::IndexFlatL2 instance, copying data over to the given GPU

GpuIndexFlatL2(std::shared_ptr<GpuResources> resources, faiss::IndexFlatL2 *index, GpuIndexFlatConfig config = GpuIndexFlatConfig())

GpuIndexFlatL2(GpuResourcesProvider *provider, int dims, GpuIndexFlatConfig config = GpuIndexFlatConfig())

Construct an empty instance that can be added to.

GpuIndexFlatL2(std::shared_ptr<GpuResources> resources, int dims, GpuIndexFlatConfig config = GpuIndexFlatConfig())

void copyFrom(faiss::IndexFlat *index)

Initialize ourselves from the given CPU index; will overwrite all data in ourselves

void copyTo(faiss::IndexFlat *index)

Copy ourselves to the given CPU index; will overwrite all data in the index instance

void copyFrom(const faiss::IndexFlat *index)

Initialize ourselves from the given CPU index; will overwrite all data in ourselves

void copyTo(faiss::IndexFlat *index) const

Copy ourselves to the given CPU index; will overwrite all data in the index instance

size_t getNumVecs() const

Returns the number of vectors we contain.

virtual void reset() override

Clears all vectors from this index.

virtual void train(idx_t n, const float *x) override

This index is not trained, so this does nothing.

virtual void add(idx_t, const float *x) override

Overrides to avoid excessive copies.

virtual void reconstruct(idx_t key, float *out) const override

Reconstruction methods; prefer the batch reconstruct as it will be more efficient

virtual void reconstruct_n(idx_t i0, idx_t num, float *out) const override

Batch reconstruction method.

virtual void reconstruct_batch(idx_t n, const idx_t *keys, float *out) const override

Batch reconstruction method.

virtual void compute_residual(const float *x, float *residual, idx_t key) const override

Compute residual.

virtual void compute_residual_n(idx_t n, const float *xs, float *residuals, const idx_t *keys) const override

Compute residual (batch mode)

inline FlatIndex *getGpuData()

For internal access.

int getDevice() const

Returns the device that this index is resident on.

std::shared_ptr<GpuResources> getResources()

Returns a reference to our GpuResources object that manages memory, stream and handle resources on the GPU

void setMinPagingSize(size_t size)

Set the minimum data size for searches (in MiB) for which we use CPU -> GPU paging

size_t getMinPagingSize() const

Returns the current minimum data size for paged searches.

virtual void add_with_ids(idx_t n, const float *x, const idx_t *ids) override

x and ids can be resident on the CPU or any GPU; copies are performed as needed Handles paged adds if the add set is too large; calls addInternal_

virtual void assign(idx_t n, const float *x, idx_t *labels, idx_t k = 1) const override

x and labels can be resident on the CPU or any GPU; copies are performed as needed

virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels, const SearchParameters *params = nullptr) const override

x, distances and labels can be resident on the CPU or any GPU; copies are performed as needed

virtual void search_and_reconstruct(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels, float *recons, const SearchParameters *params = nullptr) const override

x, distances and labels and recons can be resident on the CPU or any GPU; copies are performed as needed

virtual void range_search(idx_t n, const float *x, float radius, RangeSearchResult *result, const SearchParameters *params = nullptr) const

query n vectors of dimension d to the index.

return all vectors with distance < radius. Note that many indexes do not implement the range_search (only the k-NN search is mandatory).

Parameters:

• x – input vectors to search, size n * d

• radius – search radius

• result – result table

virtual size_t remove_ids(const IDSelector &sel)

removes IDs from the index. Not supported by all indexes. Returns the number of elements removed.

virtual DistanceComputer *get_distance_computer() const

Get a DistanceComputer (defined in AuxIndexStructures) object for this kind of index.

DistanceComputer is implemented for indexes that support random access of their vectors.

virtual size_t sa_code_size() const

size of the produced codes in bytes

virtual void sa_encode(idx_t n, const float *x, uint8_t *bytes) const

encode a set of vectors

Parameters:

• n – number of vectors

• x – input vectors, size n * d

• bytes – output encoded vectors, size n * sa_code_size()

virtual void sa_decode(idx_t n, const uint8_t *bytes, float *x) const

decode a set of vectors

Parameters:

• n – number of vectors

• bytes – input encoded vectors, size n * sa_code_size()

• x – output vectors, size n * d

virtual void merge_from(Index &otherIndex, idx_t add_id = 0)

moves the entries from another dataset to self. On output, other is empty. add_id is added to all moved ids (for sequential ids, this would be this->ntotal)

virtual void check_compatible_for_merge(const Index &otherIndex) const

check that the two indexes are compatible (ie, they are trained in the same way and have the same parameters). Otherwise throw.

Public Members

int d

vector dimension

idx_t ntotal

total nb of indexed vectors

bool verbose

verbosity level

bool is_trained

set if the Index does not require training, or if training is done already

MetricType metric_type

type of metric this index uses for search

float metric_arg

argument of the metric type

Protected Functions

void copyFrom(const faiss::Index *index)

Copy what we need from the CPU equivalent.

void copyTo(faiss::Index *index) const

Copy what we have to the CPU equivalent.

virtual bool addImplRequiresIDs_() const override

Flat index does not require IDs as there is no storage available for them

virtual void addImpl_(int n, const float *x, const idx_t *ids) override

Called from GpuIndex for add.

virtual void searchImpl_(int n, const float *x, int k, float *distances, idx_t *labels, const SearchParameters *params) const override

Called from GpuIndex for search.

Protected Attributes

const GpuIndexFlatConfig flatConfig_

Our configuration options.

std::unique_ptr<FlatIndex> data_

Holds our GPU data containing the list of vectors.

std::shared_ptr<GpuResources> resources_

Manages streams, cuBLAS handles and scratch memory for devices.

const GpuIndexConfig config_

Our configuration options.

size_t minPagedSize_

Size above which we page copies from the CPU to GPU.

class GpuIndexFlatIP : public faiss::gpu::GpuIndexFlat

#include <GpuIndexFlat.h>

Wrapper around the GPU implementation that looks like faiss::IndexFlatIP; copies over centroid data from a given faiss::IndexFlat

Public Types

using component_t = float

using distance_t = float

Public Functions

GpuIndexFlatIP(GpuResourcesProvider *provider, faiss::IndexFlatIP *index, GpuIndexFlatConfig config = GpuIndexFlatConfig())

Construct from a pre-existing faiss::IndexFlatIP instance, copying data over to the given GPU

GpuIndexFlatIP(std::shared_ptr<GpuResources> resources, faiss::IndexFlatIP *index, GpuIndexFlatConfig config = GpuIndexFlatConfig())

GpuIndexFlatIP(GpuResourcesProvider *provider, int dims, GpuIndexFlatConfig config = GpuIndexFlatConfig())

Construct an empty instance that can be added to.

GpuIndexFlatIP(std::shared_ptr<GpuResources> resources, int dims, GpuIndexFlatConfig config = GpuIndexFlatConfig())

void copyFrom(faiss::IndexFlat *index)

Initialize ourselves from the given CPU index; will overwrite all data in ourselves

void copyTo(faiss::IndexFlat *index)

Copy ourselves to the given CPU index; will overwrite all data in the index instance

void copyFrom(const faiss::IndexFlat *index)

Initialize ourselves from the given CPU index; will overwrite all data in ourselves

void copyTo(faiss::IndexFlat *index) const

Copy ourselves to the given CPU index; will overwrite all data in the index instance

size_t getNumVecs() const

Returns the number of vectors we contain.

virtual void reset() override

Clears all vectors from this index.

virtual void train(idx_t n, const float *x) override

This index is not trained, so this does nothing.

virtual void add(idx_t, const float *x) override

Overrides to avoid excessive copies.

virtual void reconstruct(idx_t key, float *out) const override

Reconstruction methods; prefer the batch reconstruct as it will be more efficient

virtual void reconstruct_n(idx_t i0, idx_t num, float *out) const override

Batch reconstruction method.

virtual void reconstruct_batch(idx_t n, const idx_t *keys, float *out) const override

Batch reconstruction method.

virtual void compute_residual(const float *x, float *residual, idx_t key) const override

Compute residual.

virtual void compute_residual_n(idx_t n, const float *xs, float *residuals, const idx_t *keys) const override

Compute residual (batch mode)

inline FlatIndex *getGpuData()

For internal access.

int getDevice() const

Returns the device that this index is resident on.

std::shared_ptr<GpuResources> getResources()

Returns a reference to our GpuResources object that manages memory, stream and handle resources on the GPU

void setMinPagingSize(size_t size)

Set the minimum data size for searches (in MiB) for which we use CPU -> GPU paging

size_t getMinPagingSize() const

Returns the current minimum data size for paged searches.

virtual void add_with_ids(idx_t n, const float *x, const idx_t *ids) override

x and ids can be resident on the CPU or any GPU; copies are performed as needed Handles paged adds if the add set is too large; calls addInternal_

virtual void assign(idx_t n, const float *x, idx_t *labels, idx_t k = 1) const override

x and labels can be resident on the CPU or any GPU; copies are performed as needed

virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels, const SearchParameters *params = nullptr) const override

x, distances and labels can be resident on the CPU or any GPU; copies are performed as needed

virtual void search_and_reconstruct(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels, float *recons, const SearchParameters *params = nullptr) const override

x, distances and labels and recons can be resident on the CPU or any GPU; copies are performed as needed

virtual void range_search(idx_t n, const float *x, float radius, RangeSearchResult *result, const SearchParameters *params = nullptr) const

query n vectors of dimension d to the index.

return all vectors with distance < radius. Note that many indexes do not implement the range_search (only the k-NN search is mandatory).

Parameters:

• x – input vectors to search, size n * d

• radius – search radius

• result – result table

virtual size_t remove_ids(const IDSelector &sel)

removes IDs from the index. Not supported by all indexes. Returns the number of elements removed.

virtual DistanceComputer *get_distance_computer() const

Get a DistanceComputer (defined in AuxIndexStructures) object for this kind of index.

DistanceComputer is implemented for indexes that support random access of their vectors.

virtual size_t sa_code_size() const

size of the produced codes in bytes

virtual void sa_encode(idx_t n, const float *x, uint8_t *bytes) const

encode a set of vectors

Parameters:

• n – number of vectors

• x – input vectors, size n * d

• bytes – output encoded vectors, size n * sa_code_size()

virtual void sa_decode(idx_t n, const uint8_t *bytes, float *x) const

decode a set of vectors

Parameters:

• n – number of vectors

• bytes – input encoded vectors, size n * sa_code_size()

• x – output vectors, size n * d

virtual void merge_from(Index &otherIndex, idx_t add_id = 0)

moves the entries from another dataset to self. On output, other is empty. add_id is added to all moved ids (for sequential ids, this would be this->ntotal)

virtual void check_compatible_for_merge(const Index &otherIndex) const

check that the two indexes are compatible (ie, they are trained in the same way and have the same parameters). Otherwise throw.

Public Members

int d

vector dimension

idx_t ntotal

total nb of indexed vectors

bool verbose

verbosity level

bool is_trained

set if the Index does not require training, or if training is done already

MetricType metric_type

type of metric this index uses for search

float metric_arg

argument of the metric type

Protected Functions

void copyFrom(const faiss::Index *index)

Copy what we need from the CPU equivalent.

void copyTo(faiss::Index *index) const

Copy what we have to the CPU equivalent.

virtual bool addImplRequiresIDs_() const override

Flat index does not require IDs as there is no storage available for them

virtual void addImpl_(int n, const float *x, const idx_t *ids) override

Called from GpuIndex for add.

virtual void searchImpl_(int n, const float *x, int k, float *distances, idx_t *labels, const SearchParameters *params) const override

Called from GpuIndex for search.

Protected Attributes

const GpuIndexFlatConfig flatConfig_

Our configuration options.

std::unique_ptr<FlatIndex> data_

Holds our GPU data containing the list of vectors.

std::shared_ptr<GpuResources> resources_

Manages streams, cuBLAS handles and scratch memory for devices.

const GpuIndexConfig config_

Our configuration options.

size_t minPagedSize_

Size above which we page copies from the CPU to GPU.

struct GpuIndexIVFConfig : public faiss::gpu::GpuIndexConfig

Subclassed by faiss::gpu::GpuIndexIVFFlatConfig, faiss::gpu::GpuIndexIVFPQConfig, faiss::gpu::GpuIndexIVFScalarQuantizerConfig

Public Functions

inline GpuIndexIVFConfig()

Public Members

IndicesOptions indicesOptions

Index storage options for the GPU.

GpuIndexFlatConfig flatConfig

Configuration for the coarse quantizer object.

int device

GPU device on which the index is resident.

MemorySpace memorySpace

What memory space to use for primary storage. On Pascal and above (CC 6+) architectures, allows GPUs to use more memory than is available on the GPU.

class GpuIndexIVF : public faiss::gpu::GpuIndex

#include <GpuIndexIVF.h>

Base class of all GPU IVF index types. This (for now) deliberately does not inherit from IndexIVF, as many of the public data members and functionality in IndexIVF is not supported in the same manner on the GPU.

Subclassed by faiss::gpu::GpuIndexIVFFlat, faiss::gpu::GpuIndexIVFPQ, faiss::gpu::GpuIndexIVFScalarQuantizer

Public Types

using component_t = float

using distance_t = float

Public Functions

GpuIndexIVF(GpuResourcesProvider *provider, int dims, faiss::MetricType metric, float metricArg, int nlist, GpuIndexIVFConfig config = GpuIndexIVFConfig())

Version that auto-constructs a flat coarse quantizer based on the desired metric

GpuIndexIVF(GpuResourcesProvider *provider, Index *coarseQuantizer, int dims, faiss::MetricType metric, float metricArg, int nlist, GpuIndexIVFConfig config = GpuIndexIVFConfig())

Version that takes a coarse quantizer instance. The GpuIndexIVF does not own the coarseQuantizer instance by default (functions like IndexIVF).

~GpuIndexIVF() override

void copyFrom(const faiss::IndexIVF *index)

Copy what we need from the CPU equivalent.

void copyTo(faiss::IndexIVF *index) const

Copy what we have to the CPU equivalent.

virtual void updateQuantizer() = 0

Should be called if the user ever changes the state of the IVF coarse quantizer manually (e.g., substitutes a new instance or changes vectors in the coarse quantizer outside the scope of training)

int getNumLists() const

Returns the number of inverted lists we’re managing.

int getListLength(int listId) const

Returns the number of vectors present in a particular inverted list.

std::vector<uint8_t> getListVectorData(int listId, bool gpuFormat = false) const

Return the encoded vector data contained in a particular inverted list, for debugging purposes. If gpuFormat is true, the data is returned as it is encoded in the GPU-side representation. Otherwise, it is converted to the CPU format. compliant format, while the native GPU format may differ.

std::vector<idx_t> getListIndices(int listId) const

Return the vector indices contained in a particular inverted list, for debugging purposes.

void setNumProbes(int nprobe)

Sets the number of list probes per query.

int getNumProbes() const

Returns our current number of list probes per query.

void search_preassigned(idx_t n, const float *x, idx_t k, const idx_t *assign, const float *centroid_dis, float *distances, idx_t *labels, bool store_pairs, const SearchParametersIVF *params = nullptr) const

Same interface as faiss::IndexIVF, in order to search a set of vectors pre-quantized by the IVF quantizer. Does not include IndexIVFStats as that can only be obtained on the host via a GPU d2h copy.

Parameters:

• n – nb of vectors to query

• x – query vectors, size nx * d

• assign – coarse quantization indices, size nx * nprobe

• centroid_dis – distances to coarse centroids, size nx * nprobe

• distance – output distances, size n * k

• labels – output labels, size n * k

• store_pairs – store inv list index + inv list offset instead in upper/lower 32 bit of result, instead of ids (used for reranking).

• params – used to override the object’s search parameters

int getDevice() const

Returns the device that this index is resident on.

std::shared_ptr<GpuResources> getResources()

Returns a reference to our GpuResources object that manages memory, stream and handle resources on the GPU

void setMinPagingSize(size_t size)

Set the minimum data size for searches (in MiB) for which we use CPU -> GPU paging

size_t getMinPagingSize() const

Returns the current minimum data size for paged searches.

virtual void add(idx_t, const float *x) override

x can be resident on the CPU or any GPU; copies are performed as needed Handles paged adds if the add set is too large; calls addInternal_

virtual void add_with_ids(idx_t n, const float *x, const idx_t *ids) override

x and ids can be resident on the CPU or any GPU; copies are performed as needed Handles paged adds if the add set is too large; calls addInternal_

virtual void assign(idx_t n, const float *x, idx_t *labels, idx_t k = 1) const override

x and labels can be resident on the CPU or any GPU; copies are performed as needed

virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels, const SearchParameters *params = nullptr) const override

x, distances and labels can be resident on the CPU or any GPU; copies are performed as needed

virtual void search_and_reconstruct(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels, float *recons, const SearchParameters *params = nullptr) const override

x, distances and labels and recons can be resident on the CPU or any GPU; copies are performed as needed

virtual void compute_residual(const float *x, float *residual, idx_t key) const override

Overridden to force GPU indices to provide their own GPU-friendly implementation

virtual void compute_residual_n(idx_t n, const float *xs, float *residuals, const idx_t *keys) const override

Overridden to force GPU indices to provide their own GPU-friendly implementation

virtual void train(idx_t n, const float *x)

Perform training on a representative set of vectors

Parameters:

• n – nb of training vectors

• x – training vecors, size n * d

virtual void range_search(idx_t n, const float *x, float radius, RangeSearchResult *result, const SearchParameters *params = nullptr) const

query n vectors of dimension d to the index.

return all vectors with distance < radius. Note that many indexes do not implement the range_search (only the k-NN search is mandatory).

Parameters:

• x – input vectors to search, size n * d

• radius – search radius

• result – result table

virtual void reset() = 0

removes all elements from the database.

virtual size_t remove_ids(const IDSelector &sel)

removes IDs from the index. Not supported by all indexes. Returns the number of elements removed.

virtual void reconstruct(idx_t key, float *recons) const

Reconstruct a stored vector (or an approximation if lossy coding)

this function may not be defined for some indexes

Parameters:

• key – id of the vector to reconstruct

• recons – reconstucted vector (size d)

virtual void reconstruct_batch(idx_t n, const idx_t *keys, float *recons) const

Reconstruct several stored vectors (or an approximation if lossy coding)

this function may not be defined for some indexes

Parameters:

• n – number of vectors to reconstruct

• keys – ids of the vectors to reconstruct (size n)

• recons – reconstucted vector (size n * d)

virtual void reconstruct_n(idx_t i0, idx_t ni, float *recons) const

Reconstruct vectors i0 to i0 + ni - 1

this function may not be defined for some indexes

Parameters:

recons – reconstucted vector (size ni * d)

virtual DistanceComputer *get_distance_computer() const

Get a DistanceComputer (defined in AuxIndexStructures) object for this kind of index.

DistanceComputer is implemented for indexes that support random access of their vectors.

virtual size_t sa_code_size() const

size of the produced codes in bytes

virtual void sa_encode(idx_t n, const float *x, uint8_t *bytes) const

encode a set of vectors

Parameters:

• n – number of vectors

• x – input vectors, size n * d

• bytes – output encoded vectors, size n * sa_code_size()

virtual void sa_decode(idx_t n, const uint8_t *bytes, float *x) const

decode a set of vectors

Parameters:

• n – number of vectors

• bytes – input encoded vectors, size n * sa_code_size()

• x – output vectors, size n * d

virtual void merge_from(Index &otherIndex, idx_t add_id = 0)

moves the entries from another dataset to self. On output, other is empty. add_id is added to all moved ids (for sequential ids, this would be this->ntotal)

virtual void check_compatible_for_merge(const Index &otherIndex) const

check that the two indexes are compatible (ie, they are trained in the same way and have the same parameters). Otherwise throw.

Public Members

ClusteringParameters cp

Exposing this like the CPU version for manipulation.

int nlist

Exposing this like the CPU version for query.

int nprobe

Exposing this like the CPU version for manipulation.

Index *quantizer

A user-pluggable coarse quantizer.

bool own_fields

Whether or not we own the coarse quantizer.

int d

vector dimension

idx_t ntotal

total nb of indexed vectors

bool verbose

verbosity level

bool is_trained

set if the Index does not require training, or if training is done already

MetricType metric_type

type of metric this index uses for search

float metric_arg

argument of the metric type

Protected Functions

void verifyIVFSettings_() const

virtual bool addImplRequiresIDs_() const override

Does addImpl_ require IDs? If so, and no IDs are provided, we will generate them sequentially based on the order in which the IDs are added

void trainQuantizer_(idx_t n, const float *x)

virtual void addImpl_(int n, const float *x, const idx_t *ids) override

Called from GpuIndex for add/add_with_ids.

virtual void searchImpl_(int n, const float *x, int k, float *distances, idx_t *labels, const SearchParameters *params) const override

Called from GpuIndex for search.

void copyFrom(const faiss::Index *index)

Copy what we need from the CPU equivalent.

void copyTo(faiss::Index *index) const

Copy what we have to the CPU equivalent.

Protected Attributes

const GpuIndexIVFConfig ivfConfig_

Our configuration options.

std::shared_ptr<IVFBase> baseIndex_

For a trained/initialized index, this is a reference to the base class.

std::shared_ptr<GpuResources> resources_

Manages streams, cuBLAS handles and scratch memory for devices.

const GpuIndexConfig config_

Our configuration options.

size_t minPagedSize_

Size above which we page copies from the CPU to GPU.

Private Functions

void init_()

Shared initialization functions.

struct GpuIndexIVFFlatConfig : public faiss::gpu::GpuIndexIVFConfig

Public Functions

inline GpuIndexIVFFlatConfig()

Public Members

bool interleavedLayout

Use the alternative memory layout for the IVF lists (currently the default)

IndicesOptions indicesOptions

Index storage options for the GPU.

GpuIndexFlatConfig flatConfig

Configuration for the coarse quantizer object.

int device

GPU device on which the index is resident.

MemorySpace memorySpace

What memory space to use for primary storage. On Pascal and above (CC 6+) architectures, allows GPUs to use more memory than is available on the GPU.

class GpuIndexIVFFlat : public faiss::gpu::GpuIndexIVF

#include <GpuIndexIVFFlat.h>

Wrapper around the GPU implementation that looks like faiss::IndexIVFFlat

Public Types

using component_t = float

using distance_t = float

Public Functions

GpuIndexIVFFlat(GpuResourcesProvider *provider, const faiss::IndexIVFFlat *index, GpuIndexIVFFlatConfig config = GpuIndexIVFFlatConfig())

Construct from a pre-existing faiss::IndexIVFFlat instance, copying data over to the given GPU, if the input index is trained.

GpuIndexIVFFlat(GpuResourcesProvider *provider, int dims, int nlist, faiss::MetricType metric = faiss::METRIC_L2, GpuIndexIVFFlatConfig config = GpuIndexIVFFlatConfig())

Constructs a new instance with an empty flat quantizer; the user provides the number of IVF lists desired.

GpuIndexIVFFlat(GpuResourcesProvider *provider, Index *coarseQuantizer, int dims, int nlist, faiss::MetricType metric = faiss::METRIC_L2, GpuIndexIVFFlatConfig config = GpuIndexIVFFlatConfig())

Constructs a new instance with a provided CPU or GPU coarse quantizer; the user provides the number of IVF lists desired.

~GpuIndexIVFFlat() override

void reserveMemory(size_t numVecs)

Reserve GPU memory in our inverted lists for this number of vectors.

void copyFrom(const faiss::IndexIVFFlat *index)

Initialize ourselves from the given CPU index; will overwrite all data in ourselves

void copyTo(faiss::IndexIVFFlat *index) const

Copy ourselves to the given CPU index; will overwrite all data in the index instance

size_t reclaimMemory()

After adding vectors, one can call this to reclaim device memory to exactly the amount needed. Returns space reclaimed in bytes

virtual void reset() override

Clears out all inverted lists, but retains the coarse centroid information

virtual void updateQuantizer() override

Should be called if the user ever changes the state of the IVF coarse quantizer manually (e.g., substitutes a new instance or changes vectors in the coarse quantizer outside the scope of training)

virtual void train(idx_t n, const float *x) override

Trains the coarse quantizer based on the given vector data.

void copyFrom(const faiss::IndexIVF *index)

Copy what we need from the CPU equivalent.

void copyTo(faiss::IndexIVF *index) const

Copy what we have to the CPU equivalent.

int getNumLists() const

Returns the number of inverted lists we’re managing.

int getListLength(int listId) const

Returns the number of vectors present in a particular inverted list.

std::vector<uint8_t> getListVectorData(int listId, bool gpuFormat = false) const

Return the encoded vector data contained in a particular inverted list, for debugging purposes. If gpuFormat is true, the data is returned as it is encoded in the GPU-side representation. Otherwise, it is converted to the CPU format. compliant format, while the native GPU format may differ.

std::vector<idx_t> getListIndices(int listId) const

Return the vector indices contained in a particular inverted list, for debugging purposes.

void setNumProbes(int nprobe)

Sets the number of list probes per query.

int getNumProbes() const

Returns our current number of list probes per query.

void search_preassigned(idx_t n, const float *x, idx_t k, const idx_t *assign, const float *centroid_dis, float *distances, idx_t *labels, bool store_pairs, const SearchParametersIVF *params = nullptr) const

Same interface as faiss::IndexIVF, in order to search a set of vectors pre-quantized by the IVF quantizer. Does not include IndexIVFStats as that can only be obtained on the host via a GPU d2h copy.

Parameters:

• n – nb of vectors to query

• x – query vectors, size nx * d

• assign – coarse quantization indices, size nx * nprobe

• centroid_dis – distances to coarse centroids, size nx * nprobe

• distance – output distances, size n * k

• labels – output labels, size n * k

• store_pairs – store inv list index + inv list offset instead in upper/lower 32 bit of result, instead of ids (used for reranking).

• params – used to override the object’s search parameters

int getDevice() const

Returns the device that this index is resident on.

std::shared_ptr<GpuResources> getResources()

Returns a reference to our GpuResources object that manages memory, stream and handle resources on the GPU

void setMinPagingSize(size_t size)

Set the minimum data size for searches (in MiB) for which we use CPU -> GPU paging

size_t getMinPagingSize() const

Returns the current minimum data size for paged searches.

virtual void add(idx_t, const float *x) override

x can be resident on the CPU or any GPU; copies are performed as needed Handles paged adds if the add set is too large; calls addInternal_

virtual void add_with_ids(idx_t n, const float *x, const idx_t *ids) override

x and ids can be resident on the CPU or any GPU; copies are performed as needed Handles paged adds if the add set is too large; calls addInternal_

virtual void assign(idx_t n, const float *x, idx_t *labels, idx_t k = 1) const override

x and labels can be resident on the CPU or any GPU; copies are performed as needed

virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels, const SearchParameters *params = nullptr) const override

x, distances and labels can be resident on the CPU or any GPU; copies are performed as needed

virtual void search_and_reconstruct(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels, float *recons, const SearchParameters *params = nullptr) const override

x, distances and labels and recons can be resident on the CPU or any GPU; copies are performed as needed

virtual void compute_residual(const float *x, float *residual, idx_t key) const override

Overridden to force GPU indices to provide their own GPU-friendly implementation

virtual void compute_residual_n(idx_t n, const float *xs, float *residuals, const idx_t *keys) const override

Overridden to force GPU indices to provide their own GPU-friendly implementation

virtual void range_search(idx_t n, const float *x, float radius, RangeSearchResult *result, const SearchParameters *params = nullptr) const

query n vectors of dimension d to the index.

return all vectors with distance < radius. Note that many indexes do not implement the range_search (only the k-NN search is mandatory).

Parameters:

• x – input vectors to search, size n * d

• radius – search radius

• result – result table

virtual size_t remove_ids(const IDSelector &sel)

removes IDs from the index. Not supported by all indexes. Returns the number of elements removed.

virtual void reconstruct(idx_t key, float *recons) const

Reconstruct a stored vector (or an approximation if lossy coding)

this function may not be defined for some indexes

Parameters:

• key – id of the vector to reconstruct

• recons – reconstucted vector (size d)

virtual void reconstruct_batch(idx_t n, const idx_t *keys, float *recons) const

Reconstruct several stored vectors (or an approximation if lossy coding)

this function may not be defined for some indexes

Parameters:

• n – number of vectors to reconstruct

• keys – ids of the vectors to reconstruct (size n)

• recons – reconstucted vector (size n * d)

virtual void reconstruct_n(idx_t i0, idx_t ni, float *recons) const

Reconstruct vectors i0 to i0 + ni - 1

this function may not be defined for some indexes

Parameters:

recons – reconstucted vector (size ni * d)

virtual DistanceComputer *get_distance_computer() const

Get a DistanceComputer (defined in AuxIndexStructures) object for this kind of index.

DistanceComputer is implemented for indexes that support random access of their vectors.

virtual size_t sa_code_size() const

size of the produced codes in bytes

virtual void sa_encode(idx_t n, const float *x, uint8_t *bytes) const

encode a set of vectors

Parameters:

• n – number of vectors

• x – input vectors, size n * d

• bytes – output encoded vectors, size n * sa_code_size()

virtual void sa_decode(idx_t n, const uint8_t *bytes, float *x) const

decode a set of vectors

Parameters:

• n – number of vectors

• bytes – input encoded vectors, size n * sa_code_size()

• x – output vectors, size n * d

virtual void merge_from(Index &otherIndex, idx_t add_id = 0)

moves the entries from another dataset to self. On output, other is empty. add_id is added to all moved ids (for sequential ids, this would be this->ntotal)

virtual void check_compatible_for_merge(const Index &otherIndex) const

check that the two indexes are compatible (ie, they are trained in the same way and have the same parameters). Otherwise throw.

Public Members

ClusteringParameters cp

Exposing this like the CPU version for manipulation.

int nlist

Exposing this like the CPU version for query.

int nprobe

Exposing this like the CPU version for manipulation.

Index *quantizer

A user-pluggable coarse quantizer.

bool own_fields

Whether or not we own the coarse quantizer.

int d

vector dimension

idx_t ntotal

total nb of indexed vectors

bool verbose

verbosity level

bool is_trained

set if the Index does not require training, or if training is done already

MetricType metric_type

type of metric this index uses for search

float metric_arg

argument of the metric type

Protected Functions

void copyFrom(const faiss::Index *index)

Copy what we need from the CPU equivalent.

void copyTo(faiss::Index *index) const

Copy what we have to the CPU equivalent.

void verifyIVFSettings_() const

virtual bool addImplRequiresIDs_() const override

Does addImpl_ require IDs? If so, and no IDs are provided, we will generate them sequentially based on the order in which the IDs are added

void trainQuantizer_(idx_t n, const float *x)

virtual void addImpl_(int n, const float *x, const idx_t *ids) override

Called from GpuIndex for add/add_with_ids.

virtual void searchImpl_(int n, const float *x, int k, float *distances, idx_t *labels, const SearchParameters *params) const override

Called from GpuIndex for search.

Protected Attributes

const GpuIndexIVFFlatConfig ivfFlatConfig_

Our configuration options.

size_t reserveMemoryVecs_

Desired inverted list memory reservation.

std::shared_ptr<IVFFlat> index_

Instance that we own; contains the inverted lists.

const GpuIndexIVFConfig ivfConfig_

Our configuration options.

std::shared_ptr<IVFBase> baseIndex_

For a trained/initialized index, this is a reference to the base class.

std::shared_ptr<GpuResources> resources_

Manages streams, cuBLAS handles and scratch memory for devices.

const GpuIndexConfig config_

Our configuration options.

size_t minPagedSize_

Size above which we page copies from the CPU to GPU.

struct GpuIndexIVFPQConfig : public faiss::gpu::GpuIndexIVFConfig

Public Functions

inline GpuIndexIVFPQConfig()

Public Members

bool useFloat16LookupTables

Whether or not float16 residual distance tables are used in the list scanning kernels. When subQuantizers * 2^bitsPerCode > 16384, this is required.

bool usePrecomputedTables

Whether or not we enable the precomputed table option for search, which can substantially increase the memory requirement.

bool interleavedLayout

Use the alternative memory layout for the IVF lists WARNING: this is a feature under development, do not use!

bool useMMCodeDistance

Use GEMM-backed computation of PQ code distances for the no precomputed table version of IVFPQ. This is for debugging purposes, it should not substantially affect the results one way for another.

Note that MM code distance is enabled automatically if one uses a number of dimensions per sub-quantizer that is not natively specialized (an odd number like 7 or so).

IndicesOptions indicesOptions

Index storage options for the GPU.

GpuIndexFlatConfig flatConfig

Configuration for the coarse quantizer object.

int device

GPU device on which the index is resident.

MemorySpace memorySpace

What memory space to use for primary storage. On Pascal and above (CC 6+) architectures, allows GPUs to use more memory than is available on the GPU.

class GpuIndexIVFPQ : public faiss::gpu::GpuIndexIVF

#include <GpuIndexIVFPQ.h>

IVFPQ index for the GPU.

Public Types

using component_t = float

using distance_t = float

Public Functions

GpuIndexIVFPQ(GpuResourcesProvider *provider, const faiss::IndexIVFPQ *index, GpuIndexIVFPQConfig config = GpuIndexIVFPQConfig())

Construct from a pre-existing faiss::IndexIVFPQ instance, copying data over to the given GPU, if the input index is trained.

GpuIndexIVFPQ(GpuResourcesProvider *provider, int dims, int nlist, int subQuantizers, int bitsPerCode, faiss::MetricType metric = faiss::METRIC_L2, GpuIndexIVFPQConfig config = GpuIndexIVFPQConfig())

Constructs a new instance with an empty flat quantizer; the user provides the number of IVF lists desired.

GpuIndexIVFPQ(GpuResourcesProvider *provider, Index *coarseQuantizer, int dims, int nlist, int subQuantizers, int bitsPerCode, faiss::MetricType metric = faiss::METRIC_L2, GpuIndexIVFPQConfig config = GpuIndexIVFPQConfig())

Constructs a new instance with a provided CPU or GPU coarse quantizer; the user provides the number of IVF lists desired.

~GpuIndexIVFPQ() override

void copyFrom(const faiss::IndexIVFPQ *index)

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 copyTo(faiss::IndexIVFPQ *index) const

Copy ourselves to the given CPU index; will overwrite all data in the index instance

void reserveMemory(size_t numVecs)

Reserve GPU memory in our inverted lists for this number of vectors.

void setPrecomputedCodes(bool enable)

Enable or disable pre-computed codes.

bool getPrecomputedCodes() const

Are pre-computed codes enabled?

int getNumSubQuantizers() const

Return the number of sub-quantizers we are using.

int getBitsPerCode() const

Return the number of bits per PQ code.

int getCentroidsPerSubQuantizer() const

Return the number of centroids per PQ code (2^bits per code)

size_t reclaimMemory()

After adding vectors, one can call this to reclaim device memory to exactly the amount needed. Returns space reclaimed in bytes

virtual void reset() override

Clears out all inverted lists, but retains the coarse and product centroid information

virtual void updateQuantizer() override

Should be called if the user ever changes the state of the IVF coarse quantizer manually (e.g., substitutes a new instance or changes vectors in the coarse quantizer outside the scope of training)

virtual void train(idx_t n, const float *x) override

Trains the coarse and product quantizer based on the given vector data.

void copyFrom(const faiss::IndexIVF *index)

Copy what we need from the CPU equivalent.

void copyTo(faiss::IndexIVF *index) const

Copy what we have to the CPU equivalent.

int getNumLists() const

Returns the number of inverted lists we’re managing.

int getListLength(int listId) const

Returns the number of vectors present in a particular inverted list.

std::vector<uint8_t> getListVectorData(int listId, bool gpuFormat = false) const

Return the encoded vector data contained in a particular inverted list, for debugging purposes. If gpuFormat is true, the data is returned as it is encoded in the GPU-side representation. Otherwise, it is converted to the CPU format. compliant format, while the native GPU format may differ.

std::vector<idx_t> getListIndices(int listId) const

Return the vector indices contained in a particular inverted list, for debugging purposes.

void setNumProbes(int nprobe)

Sets the number of list probes per query.

int getNumProbes() const

Returns our current number of list probes per query.

void search_preassigned(idx_t n, const float *x, idx_t k, const idx_t *assign, const float *centroid_dis, float *distances, idx_t *labels, bool store_pairs, const SearchParametersIVF *params = nullptr) const

Same interface as faiss::IndexIVF, in order to search a set of vectors pre-quantized by the IVF quantizer. Does not include IndexIVFStats as that can only be obtained on the host via a GPU d2h copy.

Parameters:

• n – nb of vectors to query

• x – query vectors, size nx * d

• assign – coarse quantization indices, size nx * nprobe

• centroid_dis – distances to coarse centroids, size nx * nprobe

• distance – output distances, size n * k

• labels – output labels, size n * k

• store_pairs – store inv list index + inv list offset instead in upper/lower 32 bit of result, instead of ids (used for reranking).

• params – used to override the object’s search parameters

int getDevice() const

Returns the device that this index is resident on.

std::shared_ptr<GpuResources> getResources()

Returns a reference to our GpuResources object that manages memory, stream and handle resources on the GPU

void setMinPagingSize(size_t size)

Set the minimum data size for searches (in MiB) for which we use CPU -> GPU paging

size_t getMinPagingSize() const

Returns the current minimum data size for paged searches.

virtual void add(idx_t, const float *x) override

x can be resident on the CPU or any GPU; copies are performed as needed Handles paged adds if the add set is too large; calls addInternal_

virtual void add_with_ids(idx_t n, const float *x, const idx_t *ids) override

x and ids can be resident on the CPU or any GPU; copies are performed as needed Handles paged adds if the add set is too large; calls addInternal_

virtual void assign(idx_t n, const float *x, idx_t *labels, idx_t k = 1) const override

x and labels can be resident on the CPU or any GPU; copies are performed as needed

virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels, const SearchParameters *params = nullptr) const override

x, distances and labels can be resident on the CPU or any GPU; copies are performed as needed

virtual void search_and_reconstruct(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels, float *recons, const SearchParameters *params = nullptr) const override

x, distances and labels and recons can be resident on the CPU or any GPU; copies are performed as needed

virtual void compute_residual(const float *x, float *residual, idx_t key) const override

Overridden to force GPU indices to provide their own GPU-friendly implementation

virtual void compute_residual_n(idx_t n, const float *xs, float *residuals, const idx_t *keys) const override

Overridden to force GPU indices to provide their own GPU-friendly implementation

virtual void range_search(idx_t n, const float *x, float radius, RangeSearchResult *result, const SearchParameters *params = nullptr) const

query n vectors of dimension d to the index.

return all vectors with distance < radius. Note that many indexes do not implement the range_search (only the k-NN search is mandatory).

Parameters:

• x – input vectors to search, size n * d

• radius – search radius

• result – result table

virtual size_t remove_ids(const IDSelector &sel)

removes IDs from the index. Not supported by all indexes. Returns the number of elements removed.

virtual void reconstruct(idx_t key, float *recons) const

Reconstruct a stored vector (or an approximation if lossy coding)

this function may not be defined for some indexes

Parameters:

• key – id of the vector to reconstruct

• recons – reconstucted vector (size d)

virtual void reconstruct_batch(idx_t n, const idx_t *keys, float *recons) const

Reconstruct several stored vectors (or an approximation if lossy coding)

this function may not be defined for some indexes

Parameters:

• n – number of vectors to reconstruct

• keys – ids of the vectors to reconstruct (size n)

• recons – reconstucted vector (size n * d)

virtual void reconstruct_n(idx_t i0, idx_t ni, float *recons) const

Reconstruct vectors i0 to i0 + ni - 1

this function may not be defined for some indexes

Parameters:

recons – reconstucted vector (size ni * d)

virtual DistanceComputer *get_distance_computer() const

Get a DistanceComputer (defined in AuxIndexStructures) object for this kind of index.

DistanceComputer is implemented for indexes that support random access of their vectors.

virtual size_t sa_code_size() const

size of the produced codes in bytes

virtual void sa_encode(idx_t n, const float *x, uint8_t *bytes) const

encode a set of vectors

Parameters:

• n – number of vectors

• x – input vectors, size n * d

• bytes – output encoded vectors, size n * sa_code_size()

virtual void sa_decode(idx_t n, const uint8_t *bytes, float *x) const

decode a set of vectors

Parameters:

• n – number of vectors

• bytes – input encoded vectors, size n * sa_code_size()

• x – output vectors, size n * d

virtual void merge_from(Index &otherIndex, idx_t add_id = 0)

moves the entries from another dataset to self. On output, other is empty. add_id is added to all moved ids (for sequential ids, this would be this->ntotal)

virtual void check_compatible_for_merge(const Index &otherIndex) const

check that the two indexes are compatible (ie, they are trained in the same way and have the same parameters). Otherwise throw.

Public Members

ProductQuantizer pq

Like the CPU version, we expose a publically-visible ProductQuantizer for manipulation

ClusteringParameters cp

Exposing this like the CPU version for manipulation.

int nlist

Exposing this like the CPU version for query.

int nprobe

Exposing this like the CPU version for manipulation.

Index *quantizer

A user-pluggable coarse quantizer.

bool own_fields

Whether or not we own the coarse quantizer.

int d

vector dimension

idx_t ntotal

total nb of indexed vectors

bool verbose

verbosity level

bool is_trained

set if the Index does not require training, or if training is done already

MetricType metric_type

type of metric this index uses for search

float metric_arg

argument of the metric type

Protected Functions

void verifyPQSettings_() const

Throws errors if configuration settings are improper.

void trainResidualQuantizer_(idx_t n, const float *x)

Trains the PQ quantizer based on the given vector data.

void copyFrom(const faiss::Index *index)

Copy what we need from the CPU equivalent.

void copyTo(faiss::Index *index) const

Copy what we have to the CPU equivalent.

void verifyIVFSettings_() const

virtual bool addImplRequiresIDs_() const override

Does addImpl_ require IDs? If so, and no IDs are provided, we will generate them sequentially based on the order in which the IDs are added

void trainQuantizer_(idx_t n, const float *x)

virtual void addImpl_(int n, const float *x, const idx_t *ids) override

Called from GpuIndex for add/add_with_ids.

virtual void searchImpl_(int n, const float *x, int k, float *distances, idx_t *labels, const SearchParameters *params) const override

Called from GpuIndex for search.

Protected Attributes

const GpuIndexIVFPQConfig ivfpqConfig_

Our configuration options that we were initialized with.

bool usePrecomputedTables_

Runtime override: whether or not we use precomputed tables.

int subQuantizers_

Number of sub-quantizers per encoded vector.

int bitsPerCode_

Bits per sub-quantizer code.

size_t reserveMemoryVecs_

Desired inverted list memory reservation.

std::shared_ptr<IVFPQ> index_

The product quantizer instance that we own; contains the inverted lists

const GpuIndexIVFConfig ivfConfig_

Our configuration options.

std::shared_ptr<IVFBase> baseIndex_

For a trained/initialized index, this is a reference to the base class.

std::shared_ptr<GpuResources> resources_

Manages streams, cuBLAS handles and scratch memory for devices.

const GpuIndexConfig config_

Our configuration options.

size_t minPagedSize_

Size above which we page copies from the CPU to GPU.

struct GpuIndexIVFScalarQuantizerConfig : public faiss::gpu::GpuIndexIVFConfig

Public Functions

inline GpuIndexIVFScalarQuantizerConfig()

Public Members

bool interleavedLayout

Use the alternative memory layout for the IVF lists (currently the default)

IndicesOptions indicesOptions

Index storage options for the GPU.

GpuIndexFlatConfig flatConfig

Configuration for the coarse quantizer object.

int device

GPU device on which the index is resident.

MemorySpace memorySpace

What memory space to use for primary storage. On Pascal and above (CC 6+) architectures, allows GPUs to use more memory than is available on the GPU.

class GpuIndexIVFScalarQuantizer : public faiss::gpu::GpuIndexIVF

#include <GpuIndexIVFScalarQuantizer.h>

Wrapper around the GPU implementation that looks like faiss::IndexIVFScalarQuantizer

Public Types

using component_t = float

using distance_t = float

Public Functions

GpuIndexIVFScalarQuantizer(GpuResourcesProvider *provider, const faiss::IndexIVFScalarQuantizer *index, GpuIndexIVFScalarQuantizerConfig config = GpuIndexIVFScalarQuantizerConfig())

Construct from a pre-existing faiss::IndexIVFScalarQuantizer instance, copying data over to the given GPU, if the input index is trained.

GpuIndexIVFScalarQuantizer(GpuResourcesProvider *provider, int dims, int nlist, faiss::ScalarQuantizer::QuantizerType qtype, faiss::MetricType metric = MetricType::METRIC_L2, bool encodeResidual = true, GpuIndexIVFScalarQuantizerConfig config = GpuIndexIVFScalarQuantizerConfig())

Constructs a new instance with an empty flat quantizer; the user provides the number of IVF lists desired.

GpuIndexIVFScalarQuantizer(GpuResourcesProvider *provider, Index *coarseQuantizer, int dims, int nlist, faiss::ScalarQuantizer::QuantizerType qtype, faiss::MetricType metric = MetricType::METRIC_L2, bool encodeResidual = true, GpuIndexIVFScalarQuantizerConfig config = GpuIndexIVFScalarQuantizerConfig())

Constructs a new instance with a provided CPU or GPU coarse quantizer; the user provides the number of IVF lists desired.

~GpuIndexIVFScalarQuantizer() override

void reserveMemory(size_t numVecs)

Reserve GPU memory in our inverted lists for this number of vectors.

void copyFrom(const faiss::IndexIVFScalarQuantizer *index)

Initialize ourselves from the given CPU index; will overwrite all data in ourselves

void copyTo(faiss::IndexIVFScalarQuantizer *index) const

Copy ourselves to the given CPU index; will overwrite all data in the index instance

size_t reclaimMemory()

After adding vectors, one can call this to reclaim device memory to exactly the amount needed. Returns space reclaimed in bytes

virtual void reset() override

Clears out all inverted lists, but retains the coarse and scalar quantizer information

virtual void updateQuantizer() override

Should be called if the user ever changes the state of the IVF coarse quantizer manually (e.g., substitutes a new instance or changes vectors in the coarse quantizer outside the scope of training)

virtual void train(idx_t n, const float *x) override

Trains the coarse and scalar quantizer based on the given vector data.

void copyFrom(const faiss::IndexIVF *index)

Copy what we need from the CPU equivalent.

void copyTo(faiss::IndexIVF *index) const

Copy what we have to the CPU equivalent.

int getNumLists() const

Returns the number of inverted lists we’re managing.

int getListLength(int listId) const

Returns the number of vectors present in a particular inverted list.

std::vector<uint8_t> getListVectorData(int listId, bool gpuFormat = false) const

Return the encoded vector data contained in a particular inverted list, for debugging purposes. If gpuFormat is true, the data is returned as it is encoded in the GPU-side representation. Otherwise, it is converted to the CPU format. compliant format, while the native GPU format may differ.

std::vector<idx_t> getListIndices(int listId) const

Return the vector indices contained in a particular inverted list, for debugging purposes.

void setNumProbes(int nprobe)

Sets the number of list probes per query.

int getNumProbes() const

Returns our current number of list probes per query.

void search_preassigned(idx_t n, const float *x, idx_t k, const idx_t *assign, const float *centroid_dis, float *distances, idx_t *labels, bool store_pairs, const SearchParametersIVF *params = nullptr) const

Same interface as faiss::IndexIVF, in order to search a set of vectors pre-quantized by the IVF quantizer. Does not include IndexIVFStats as that can only be obtained on the host via a GPU d2h copy.

Parameters:

• n – nb of vectors to query

• x – query vectors, size nx * d

• assign – coarse quantization indices, size nx * nprobe

• centroid_dis – distances to coarse centroids, size nx * nprobe

• distance – output distances, size n * k

• labels – output labels, size n * k

• store_pairs – store inv list index + inv list offset instead in upper/lower 32 bit of result, instead of ids (used for reranking).

• params – used to override the object’s search parameters

int getDevice() const

Returns the device that this index is resident on.

std::shared_ptr<GpuResources> getResources()

Returns a reference to our GpuResources object that manages memory, stream and handle resources on the GPU

void setMinPagingSize(size_t size)

Set the minimum data size for searches (in MiB) for which we use CPU -> GPU paging

size_t getMinPagingSize() const

Returns the current minimum data size for paged searches.

virtual void add(idx_t, const float *x) override

x can be resident on the CPU or any GPU; copies are performed as needed Handles paged adds if the add set is too large; calls addInternal_

virtual void add_with_ids(idx_t n, const float *x, const idx_t *ids) override

x and ids can be resident on the CPU or any GPU; copies are performed as needed Handles paged adds if the add set is too large; calls addInternal_

virtual void assign(idx_t n, const float *x, idx_t *labels, idx_t k = 1) const override

x and labels can be resident on the CPU or any GPU; copies are performed as needed

virtual void search(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels, const SearchParameters *params = nullptr) const override

x, distances and labels can be resident on the CPU or any GPU; copies are performed as needed

virtual void search_and_reconstruct(idx_t n, const float *x, idx_t k, float *distances, idx_t *labels, float *recons, const SearchParameters *params = nullptr) const override

x, distances and labels and recons can be resident on the CPU or any GPU; copies are performed as needed

virtual void compute_residual(const float *x, float *residual, idx_t key) const override

Overridden to force GPU indices to provide their own GPU-friendly implementation

virtual void compute_residual_n(idx_t n, const float *xs, float *residuals, const idx_t *keys) const override

Overridden to force GPU indices to provide their own GPU-friendly implementation

virtual void range_search(idx_t n, const float *x, float radius, RangeSearchResult *result, const SearchParameters *params = nullptr) const

query n vectors of dimension d to the index.

return all vectors with distance < radius. Note that many indexes do not implement the range_search (only the k-NN search is mandatory).

Parameters:

• x – input vectors to search, size n * d

• radius – search radius

• result – result table

virtual size_t remove_ids(const IDSelector &sel)

removes IDs from the index. Not supported by all indexes. Returns the number of elements removed.

virtual void reconstruct(idx_t key, float *recons) const

Reconstruct a stored vector (or an approximation if lossy coding)

this function may not be defined for some indexes

Parameters:

• key – id of the vector to reconstruct

• recons – reconstucted vector (size d)

virtual void reconstruct_batch(idx_t n, const idx_t *keys, float *recons) const

Reconstruct several stored vectors (or an approximation if lossy coding)

this function may not be defined for some indexes

Parameters:

• n – number of vectors to reconstruct

• keys – ids of the vectors to reconstruct (size n)

• recons – reconstucted vector (size n * d)

virtual void reconstruct_n(idx_t i0, idx_t ni, float *recons) const

Reconstruct vectors i0 to i0 + ni - 1

this function may not be defined for some indexes

Parameters:

recons – reconstucted vector (size ni * d)

virtual DistanceComputer *get_distance_computer() const

Get a DistanceComputer (defined in AuxIndexStructures) object for this kind of index.

DistanceComputer is implemented for indexes that support random access of their vectors.

virtual size_t sa_code_size() const

size of the produced codes in bytes

virtual void sa_encode(idx_t n, const float *x, uint8_t *bytes) const

encode a set of vectors

Parameters:

• n – number of vectors

• x – input vectors, size n * d

• bytes – output encoded vectors, size n * sa_code_size()

virtual void sa_decode(idx_t n, const uint8_t *bytes, float *x) const

decode a set of vectors

Parameters:

• n – number of vectors

• bytes – input encoded vectors, size n * sa_code_size()

• x – output vectors, size n * d

virtual void merge_from(Index &otherIndex, idx_t add_id = 0)

moves the entries from another dataset to self. On output, other is empty. add_id is added to all moved ids (for sequential ids, this would be this->ntotal)

virtual void check_compatible_for_merge(const Index &otherIndex) const

check that the two indexes are compatible (ie, they are trained in the same way and have the same parameters). Otherwise throw.

Public Members

faiss::ScalarQuantizer sq

Exposed like the CPU version.

bool by_residual

Exposed like the CPU version.

ClusteringParameters cp

Exposing this like the CPU version for manipulation.

int nlist

Exposing this like the CPU version for query.

int nprobe

Exposing this like the CPU version for manipulation.

Index *quantizer

A user-pluggable coarse quantizer.

bool own_fields

Whether or not we own the coarse quantizer.

int d

vector dimension

idx_t ntotal

total nb of indexed vectors

bool verbose

verbosity level

bool is_trained

set if the Index does not require training, or if training is done already

MetricType metric_type

type of metric this index uses for search

float metric_arg

argument of the metric type

Protected Functions

void verifySQSettings_() const

Validates index SQ parameters.

void trainResiduals_(idx_t n, const float *x)

Called from train to handle SQ residual training.

void copyFrom(const faiss::Index *index)

Copy what we need from the CPU equivalent.

void copyTo(faiss::Index *index) const

Copy what we have to the CPU equivalent.

void verifyIVFSettings_() const

virtual bool addImplRequiresIDs_() const override

Does addImpl_ require IDs? If so, and no IDs are provided, we will generate them sequentially based on the order in which the IDs are added

void trainQuantizer_(idx_t n, const float *x)

virtual void addImpl_(int n, const float *x, const idx_t *ids) override

Called from GpuIndex for add/add_with_ids.

virtual void searchImpl_(int n, const float *x, int k, float *distances, idx_t *labels, const SearchParameters *params) const override

Called from GpuIndex for search.

Protected Attributes

const GpuIndexIVFScalarQuantizerConfig ivfSQConfig_

Our configuration options.