All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
GpuIndexFlat.h
/**
* Copyright (c) 2015-present, Facebook, Inc.
* All rights reserved.
*
* This source code is licensed under the CC-by-NC license found in the
* LICENSE file in the root directory of this source tree.
*/
// Copyright 2004-present Facebook. All Rights Reserved.
#pragma once
#include "GpuIndex.h"
namespace faiss {
struct IndexFlat;
struct IndexFlatL2;
struct IndexFlatIP;
}
namespace faiss { namespace gpu {
struct FlatIndex;
struct GpuIndexFlatConfig {
inline GpuIndexFlatConfig()
: device(0),
useFloat16(false),
storeTransposed(false) {
}
int device;
bool useFloat16;
bool storeTransposed;
};
/// Wrapper around the GPU implementation that looks like
/// faiss::IndexFlat; copies over centroid data from a given
/// faiss::IndexFlat
class GpuIndexFlat : public GpuIndex {
public:
/// Construct from a pre-existing faiss::IndexFlat instance, copying
/// data over to the given GPU
GpuIndexFlat(GpuResources* resources,
const faiss::IndexFlat* index,
GpuIndexFlatConfig config = GpuIndexFlatConfig());
/// Construct an empty instance that can be added to
GpuIndexFlat(GpuResources* resources,
int dims,
faiss::MetricType metric,
GpuIndexFlatConfig config = GpuIndexFlatConfig());
virtual ~GpuIndexFlat();
/// Set the minimum data size for searches (in MiB) for which we use
/// CPU -> GPU paging
void setMinPagingSize(size_t size);
/// Returns the current minimum data size for paged searches
size_t getMinPagingSize() const;
/// Do we store vectors and perform math in float16?
bool getUseFloat16() const;
/// Initialize ourselves from the given CPU index; will overwrite
/// all data in ourselves
void copyFrom(const faiss::IndexFlat* index);
/// Copy ourselves to the given CPU index; will overwrite all data
/// in the index instance
void copyTo(faiss::IndexFlat* index) const;
/// Returns the number of vectors we contain
size_t getNumVecs() const;
/// Clears all vectors from this index
virtual void reset();
/// This index is not trained, so this does nothing
virtual void train(Index::idx_t n, const float* x);
/// `x`, `distances` and `labels` can be resident on the CPU or any
/// GPU; copies are performed as needed
/// We have our own implementation here which handles CPU async
/// copies; searchImpl_ is not called
/// FIXME: move paged impl into GpuIndex
virtual void search(faiss::Index::idx_t n,
const float* x,
faiss::Index::idx_t k,
float* distances,
faiss::Index::idx_t* labels) const;
/// Reconstruction methods; prefer the batch reconstruct as it will
/// be more efficient
virtual void reconstruct(faiss::Index::idx_t key, float* out) const;
/// Batch reconstruction method
virtual void reconstruct_n(faiss::Index::idx_t i0,
faiss::Index::idx_t num,
float* out) const;
virtual void set_typename();
/// For internal access
inline FlatIndex* getGpuData() { return data_; }
protected:
/// Called from GpuIndex for add
virtual void addImpl_(faiss::Index::idx_t n,
const float* x,
const faiss::Index::idx_t* ids);
/// Should not be called (we have our own implementation)
virtual void searchImpl_(faiss::Index::idx_t n,
const float* x,
faiss::Index::idx_t k,
float* distances,
faiss::Index::idx_t* labels) const;
/// Called from search when the input data is on the CPU;
/// potentially allows for pinned memory usage
void searchFromCpuPaged_(int n,
const float* x,
int k,
float* outDistancesData,
int* outIndicesData) const;
void searchNonPaged_(int n,
const float* x,
int k,
float* outDistancesData,
int* outIndicesData) const;
protected:
/// Size above which we page copies from the CPU to GPU
size_t minPagedSize_;
const GpuIndexFlatConfig config_;
/// Holds our GPU data containing the list of vectors
FlatIndex* data_;
};
/// Wrapper around the GPU implementation that looks like
/// faiss::IndexFlatL2; copies over centroid data from a given
/// faiss::IndexFlat
class GpuIndexFlatL2 : public GpuIndexFlat {
public:
/// Construct from a pre-existing faiss::IndexFlatL2 instance, copying
/// data over to the given GPU
GpuIndexFlatL2(GpuResources* resources,
faiss::IndexFlatL2* index,
GpuIndexFlatConfig config = GpuIndexFlatConfig());
/// Construct an empty instance that can be added to
GpuIndexFlatL2(GpuResources* resources,
int dims,
GpuIndexFlatConfig config = GpuIndexFlatConfig());
/// Initialize ourselves from the given CPU index; will overwrite
/// all data in ourselves
void copyFrom(faiss::IndexFlatL2* index);
/// Copy ourselves to the given CPU index; will overwrite all data
/// in the index instance
void copyTo(faiss::IndexFlatL2* index);
};
/// Wrapper around the GPU implementation that looks like
/// faiss::IndexFlatIP; copies over centroid data from a given
/// faiss::IndexFlat
class GpuIndexFlatIP : public GpuIndexFlat {
public:
/// Construct from a pre-existing faiss::IndexFlatIP instance, copying
/// data over to the given GPU
GpuIndexFlatIP(GpuResources* resources,
faiss::IndexFlatIP* index,
GpuIndexFlatConfig config = GpuIndexFlatConfig());
/// Construct an empty instance that can be added to
GpuIndexFlatIP(GpuResources* resources,
int dims,
GpuIndexFlatConfig config = GpuIndexFlatConfig());
/// Initialize ourselves from the given CPU index; will overwrite
/// all data in ourselves
void copyFrom(faiss::IndexFlatIP* index);
/// Copy ourselves to the given CPU index; will overwrite all data
/// in the index instance
void copyTo(faiss::IndexFlatIP* index);
};
} } // namespace
faiss::IndexFlat
Definition: IndexFlat.h:24
faiss::gpu::GpuIndexFlatL2::copyFrom
void copyFrom(faiss::IndexFlatL2 *index)
Definition: GpuIndexFlat.cu:519
faiss::gpu::GpuIndexFlat::copyTo
void copyTo(faiss::IndexFlat *index) const
Definition: GpuIndexFlat.cu:111
faiss::gpu::GpuIndexFlat::searchImpl_
virtual void searchImpl_(faiss::Index::idx_t n, const float *x, faiss::Index::idx_t k, float *distances, faiss::Index::idx_t *labels) const
Should not be called (we have our own implementation)
Definition: GpuIndexFlat.cu:255
faiss::gpu::GpuIndexFlat::getMinPagingSize
size_t getMinPagingSize() const
Returns the current minimum data size for paged searches.
Definition: GpuIndexFlat.cu:74
faiss::gpu::FlatIndex
Holder of GPU resources for a particular flat index.
Definition: FlatIndex.cuh:23
faiss::gpu::GpuIndexFlat::getGpuData
FlatIndex * getGpuData()
For internal access.
Definition: GpuIndexFlat.h:109
faiss::gpu::GpuIndexFlatL2::copyTo
void copyTo(faiss::IndexFlatL2 *index)
Definition: GpuIndexFlat.cu:524
faiss::IndexFlatL2
Definition: IndexFlat.h:81
faiss::gpu::GpuIndexFlat::searchFromCpuPaged_
void searchFromCpuPaged_(int n, const float *x, int k, float *outDistancesData, int *outIndicesData) const
Definition: GpuIndexFlat.cu:286
faiss::gpu::GpuIndexFlat::getNumVecs
size_t getNumVecs() const
Returns the number of vectors we contain.
Definition: GpuIndexFlat.cu:134
faiss::gpu::GpuResources
Definition: GpuResources.h:24
faiss::gpu::GpuIndexFlat::train
virtual void train(Index::idx_t n, const float *x)
This index is not trained, so this does nothing.
Definition: GpuIndexFlat.cu:148
faiss::gpu::GpuIndexFlat::GpuIndexFlat
GpuIndexFlat(GpuResources *resources, const faiss::IndexFlat *index, GpuIndexFlatConfig config=GpuIndexFlatConfig())
Definition: GpuIndexFlat.cu:33
faiss::gpu::GpuIndexFlat::addImpl_
virtual void addImpl_(faiss::Index::idx_t n, const float *x, const faiss::Index::idx_t *ids)
Called from GpuIndex for add.
Definition: GpuIndexFlat.cu:153
faiss::IndexFlatIP
Definition: IndexFlat.h:75
faiss::gpu::GpuIndexFlat::setMinPagingSize
void setMinPagingSize(size_t size)
Definition: GpuIndexFlat.cu:69
faiss::gpu::GpuIndex
Definition: GpuIndex.h:20
faiss::gpu::GpuIndexFlatL2::GpuIndexFlatL2
GpuIndexFlatL2(GpuResources *resources, faiss::IndexFlatL2 *index, GpuIndexFlatConfig config=GpuIndexFlatConfig())
Definition: GpuIndexFlat.cu:506
faiss::gpu::GpuIndexFlat::getUseFloat16
bool getUseFloat16() const
Do we store vectors and perform math in float16?
Definition: GpuIndexFlat.cu:79
faiss::gpu::GpuIndexFlat::reset
virtual void reset()
Clears all vectors from this index.
Definition: GpuIndexFlat.cu:139
faiss::gpu::GpuIndexFlatConfig
Definition: GpuIndexFlat.h:28
faiss::gpu::GpuIndexFlatIP::copyTo
void copyTo(faiss::IndexFlatIP *index)
Definition: GpuIndexFlat.cu:550
faiss::gpu::GpuIndexFlatIP::GpuIndexFlatIP
GpuIndexFlatIP(GpuResources *resources, faiss::IndexFlatIP *index, GpuIndexFlatConfig config=GpuIndexFlatConfig())
Definition: GpuIndexFlat.cu:532
faiss::Index::idx_t
long idx_t
all indices are this type
Definition: Index.h:64
faiss::gpu::GpuIndexFlat::search
virtual void search(faiss::Index::idx_t n, const float *x, faiss::Index::idx_t k, float *distances, faiss::Index::idx_t *labels) const
Definition: GpuIndexFlat.cu:176
faiss::gpu::GpuIndexFlatIP
Definition: GpuIndexFlat.h:176
faiss::gpu::GpuIndexFlat::copyFrom
void copyFrom(const faiss::IndexFlat *index)
Definition: GpuIndexFlat.cu:84
faiss::gpu::GpuIndexFlat::minPagedSize_
size_t minPagedSize_
Size above which we page copies from the CPU to GPU.
Definition: GpuIndexFlat.h:140
faiss::gpu::GpuIndexFlatL2
Definition: GpuIndexFlat.h:151
faiss::gpu::GpuIndexFlat::reconstruct
virtual void reconstruct(faiss::Index::idx_t key, float *out) const
Definition: GpuIndexFlat.cu:458
faiss::gpu::GpuIndexFlat
Definition: GpuIndexFlat.h:43
faiss::gpu::GpuIndexFlatIP::copyFrom
void copyFrom(faiss::IndexFlatIP *index)
Definition: GpuIndexFlat.cu:545
faiss::gpu::GpuIndexFlat::reconstruct_n
virtual void reconstruct_n(faiss::Index::idx_t i0, faiss::Index::idx_t num, float *out) const
Batch reconstruction method.
Definition: GpuIndexFlat.cu:475
faiss::MetricType
MetricType
Some algorithms support both an inner product vetsion and a L2 search version.
Definition: Index.h:44
faiss::gpu::GpuIndexFlat::data_
FlatIndex * data_
Holds our GPU data containing the list of vectors.
Definition: GpuIndexFlat.h:145