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 BSD+Patents 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 : public GpuIndexConfig {
inline GpuIndexFlatConfig()
: useFloat16(false),
useFloat16Accumulator(false),
storeTransposed(false) {
}
/// Whether or not data is stored as float16
bool useFloat16;
/// Whether or not all math is performed in float16, if useFloat16 is
/// specified. If true, we use cublasHgemm, supported only on CC
/// 5.3+. Otherwise, we use cublasSgemmEx.
bool useFloat16Accumulator;
/// Whether or not data is stored (transparently) in a transposed
/// layout, enabling use of the NN GEMM call, which is ~10% faster.
/// This will improve the speed of the flat index, but will
/// substantially slow down any add() calls made, as all data must
/// be transposed, and will increase storage requirements (we store
/// data in both transposed and non-transposed layouts).
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());
~GpuIndexFlat() override;
/// 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;
/// 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
void reset() override;
/// This index is not trained, so this does nothing
void train(Index::idx_t n, const float* x) override;
/// Overrides to avoid excessive copies
void add(faiss::Index::idx_t, const float* x) override;
/// `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
void search(
faiss::Index::idx_t n,
const float* x,
faiss::Index::idx_t k,
float* distances,
faiss::Index::idx_t* labels) const override;
/// Reconstruction methods; prefer the batch reconstruct as it will
/// be more efficient
void reconstruct(faiss::Index::idx_t key, float* out) const override;
/// Batch reconstruction method
void reconstruct_n(
faiss::Index::idx_t i0,
faiss::Index::idx_t num,
float* out) const override;
/// For internal access
inline FlatIndex* getGpuData() { return data_; }
protected:
/// Called from GpuIndex for add
void addImpl_(
faiss::Index::idx_t n,
const float* x,
const faiss::Index::idx_t* ids) override;
/// Should not be called (we have our own implementation)
void searchImpl_(
faiss::Index::idx_t n,
const float* x,
faiss::Index::idx_t k,
float* distances,
faiss::Index::idx_t* labels) const override;
/// Called from 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;
private:
/// Checks user settings for consistency
void verifySettings_() const;
protected:
/// Our config object
const GpuIndexFlatConfig config_;
/// Size above which we page copies from the CPU to GPU
size_t minPagedSize_;
/// 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:23
faiss::gpu::GpuIndexFlatConfig::useFloat16Accumulator
bool useFloat16Accumulator
Definition: GpuIndexFlat.h:40
faiss::gpu::GpuIndexFlatConfig::storeTransposed
bool storeTransposed
Definition: GpuIndexFlat.h:48
faiss::gpu::GpuIndexFlatL2::copyFrom
void copyFrom(faiss::IndexFlatL2 *index)
Definition: GpuIndexFlat.cu:562
faiss::gpu::GpuIndexFlat::copyTo
void copyTo(faiss::IndexFlat *index) const
Definition: GpuIndexFlat.cu:120
faiss::gpu::GpuIndexFlat::reconstruct_n
void reconstruct_n(faiss::Index::idx_t i0, faiss::Index::idx_t num, float *out) const override
Batch reconstruction method.
Definition: GpuIndexFlat.cu:508
faiss::gpu::GpuIndexFlat::getMinPagingSize
size_t getMinPagingSize() const
Returns the current minimum data size for paged searches.
Definition: GpuIndexFlat.cu:82
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:120
faiss::gpu::GpuIndexFlatL2::copyTo
void copyTo(faiss::IndexFlatL2 *index)
Definition: GpuIndexFlat.cu:567
faiss::IndexFlatL2
Definition: IndexFlat.h:79
faiss::gpu::GpuIndexFlat::searchFromCpuPaged_
void searchFromCpuPaged_(int n, const float *x, int k, float *outDistancesData, int *outIndicesData) const
Definition: GpuIndexFlat.cu:319
faiss::gpu::GpuIndexFlat::addImpl_
void addImpl_(faiss::Index::idx_t n, const float *x, const faiss::Index::idx_t *ids) override
Called from GpuIndex for add.
Definition: GpuIndexFlat.cu:180
faiss::gpu::GpuIndexFlat::searchImpl_
void searchImpl_(faiss::Index::idx_t n, const float *x, faiss::Index::idx_t k, float *distances, faiss::Index::idx_t *labels) const override
Should not be called (we have our own implementation)
Definition: GpuIndexFlat.cu:288
faiss::gpu::GpuIndexFlat::getNumVecs
size_t getNumVecs() const
Returns the number of vectors we contain.
Definition: GpuIndexFlat.cu:143
faiss::gpu::GpuResources
Definition: GpuResources.h:23
faiss::gpu::GpuIndexFlat::GpuIndexFlat
GpuIndexFlat(GpuResources *resources, const faiss::IndexFlat *index, GpuIndexFlatConfig config=GpuIndexFlatConfig())
Definition: GpuIndexFlat.cu:33
faiss::IndexFlatIP
Definition: IndexFlat.h:73
faiss::gpu::GpuIndexFlat::setMinPagingSize
void setMinPagingSize(size_t size)
Definition: GpuIndexFlat.cu:77
faiss::gpu::GpuIndexFlatConfig::useFloat16
bool useFloat16
Whether or not data is stored as float16.
Definition: GpuIndexFlat.h:35
faiss::gpu::GpuIndex
Definition: GpuIndex.h:35
faiss::gpu::GpuIndexFlatL2::GpuIndexFlatL2
GpuIndexFlatL2(GpuResources *resources, faiss::IndexFlatL2 *index, GpuIndexFlatConfig config=GpuIndexFlatConfig())
Definition: GpuIndexFlat.cu:549
faiss::gpu::GpuIndexFlatConfig
Definition: GpuIndexFlat.h:27
faiss::gpu::GpuIndexFlat::reconstruct
void reconstruct(faiss::Index::idx_t key, float *out) const override
Definition: GpuIndexFlat.cu:491
faiss::gpu::GpuIndexFlatIP::copyTo
void copyTo(faiss::IndexFlatIP *index)
Definition: GpuIndexFlat.cu:593
faiss::gpu::GpuIndexFlatIP::GpuIndexFlatIP
GpuIndexFlatIP(GpuResources *resources, faiss::IndexFlatIP *index, GpuIndexFlatConfig config=GpuIndexFlatConfig())
Definition: GpuIndexFlat.cu:575
faiss::Index::idx_t
long idx_t
all indices are this type
Definition: Index.h:62
faiss::gpu::GpuIndexFlatIP
Definition: GpuIndexFlat.h:194
faiss::gpu::GpuIndexFlat::copyFrom
void copyFrom(const faiss::IndexFlat *index)
Definition: GpuIndexFlat.cu:87
faiss::gpu::GpuIndexFlat::minPagedSize_
size_t minPagedSize_
Size above which we page copies from the CPU to GPU.
Definition: GpuIndexFlat.h:160
faiss::gpu::GpuIndexFlat::config_
const GpuIndexFlatConfig config_
Our config object.
Definition: GpuIndexFlat.h:157
faiss::gpu::GpuIndexFlatL2
Definition: GpuIndexFlat.h:169
faiss::gpu::GpuIndexConfig
Definition: GpuIndex.h:20
faiss::gpu::GpuIndexFlat::add
void add(faiss::Index::idx_t, const float *x) override
Overrides to avoid excessive copies.
Definition: GpuIndexFlat.cu:162
faiss::gpu::GpuIndexFlat
Definition: GpuIndexFlat.h:54
faiss::gpu::GpuIndexFlat::reset
void reset() override
Clears all vectors from this index.
Definition: GpuIndexFlat.cu:148
faiss::gpu::GpuIndexFlatIP::copyFrom
void copyFrom(faiss::IndexFlatIP *index)
Definition: GpuIndexFlat.cu:588
faiss::gpu::GpuIndexFlat::search
void search(faiss::Index::idx_t n, const float *x, faiss::Index::idx_t k, float *distances, faiss::Index::idx_t *labels) const override
Definition: GpuIndexFlat.cu:205
faiss::gpu::GpuIndexFlat::train
void train(Index::idx_t n, const float *x) override
This index is not trained, so this does nothing.
Definition: GpuIndexFlat.cu:157
faiss::MetricType
MetricType
Some algorithms support both an inner product vetsion and a L2 search version.
Definition: Index.h:43
faiss::gpu::GpuIndexFlat::data_
FlatIndex * data_
Holds our GPU data containing the list of vectors.
Definition: GpuIndexFlat.h:163