Class faiss::gpu::GpuIndex

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:
  • n – number of vectors

  • 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:
  • i0 – index of the first vector in the sequence

  • ni – number of vectors in the sequence

  • 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_(idx_t 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_(idx_t 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_(idx_t 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_(idx_t n, const float *x, const idx_t *ids)

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

void searchNonPaged_(idx_t 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_(idx_t 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