Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
GpuIndexFlat.h
1 /**
2 * Copyright (c) Facebook, Inc. and its affiliates.
3 *
4 * This source code is licensed under the MIT license found in the
5 * LICENSE file in the root directory of this source tree.
6 */
7
8
9 #pragma once
10
11 #include "GpuIndex.h"
12
13 namespace faiss {
14
15 struct IndexFlat;
16 struct IndexFlatL2;
17 struct IndexFlatIP;
18
19 }
20
21 namespace faiss { namespace gpu {
22
23 struct FlatIndex;
24
25 struct GpuIndexFlatConfig : public GpuIndexConfig {
26 inline GpuIndexFlatConfig()
27 : useFloat16(false),
28 useFloat16Accumulator(false),
29 storeTransposed(false) {
30 }
31
32 /// Whether or not data is stored as float16
33 bool useFloat16;
34
35 /// Whether or not all math is performed in float16, if useFloat16 is
36 /// specified. If true, we use cublasHgemm, supported only on CC
37 /// 5.3+. Otherwise, we use cublasSgemmEx.
38 bool useFloat16Accumulator;
39
40 /// Whether or not data is stored (transparently) in a transposed
41 /// layout, enabling use of the NN GEMM call, which is ~10% faster.
42 /// This will improve the speed of the flat index, but will
43 /// substantially slow down any add() calls made, as all data must
44 /// be transposed, and will increase storage requirements (we store
45 /// data in both transposed and non-transposed layouts).
46 bool storeTransposed;
47 };
48
49 /// Wrapper around the GPU implementation that looks like
50 /// faiss::IndexFlat; copies over centroid data from a given
51 /// faiss::IndexFlat
52 class GpuIndexFlat : public GpuIndex {
53 public:
54 /// Construct from a pre-existing faiss::IndexFlat instance, copying
55 /// data over to the given GPU
56 GpuIndexFlat(GpuResources* resources,
57 const faiss::IndexFlat* index,
58 GpuIndexFlatConfig config = GpuIndexFlatConfig());
59
60 /// Construct an empty instance that can be added to
61 GpuIndexFlat(GpuResources* resources,
62 int dims,
63 faiss::MetricType metric,
64 GpuIndexFlatConfig config = GpuIndexFlatConfig());
65
66 ~GpuIndexFlat() override;
67
68 /// Initialize ourselves from the given CPU index; will overwrite
69 /// all data in ourselves
70 void copyFrom(const faiss::IndexFlat* index);
71
72 /// Copy ourselves to the given CPU index; will overwrite all data
73 /// in the index instance
74 void copyTo(faiss::IndexFlat* index) const;
75
76 /// Returns the number of vectors we contain
77 size_t getNumVecs() const;
78
79 /// Clears all vectors from this index
80 void reset() override;
81
82 /// This index is not trained, so this does nothing
83 void train(Index::idx_t n, const float* x) override;
84
85 /// Overrides to avoid excessive copies
86 void add(faiss::Index::idx_t, const float* x) override;
87
88 /// Reconstruction methods; prefer the batch reconstruct as it will
89 /// be more efficient
90 void reconstruct(faiss::Index::idx_t key, float* out) const override;
91
92 /// Batch reconstruction method
93 void reconstruct_n(
94 faiss::Index::idx_t i0,
95 faiss::Index::idx_t num,
96 float* out) const override;
97
98 /// For internal access
99 inline FlatIndex* getGpuData() { return data_; }
100
101 protected:
102 /// Flat index does not require IDs as there is no storage available for them
103 bool addImplRequiresIDs_() const override;
104
105 /// Called from GpuIndex for add
106 void addImpl_(int n,
107 const float* x,
108 const Index::idx_t* ids) override;
109
110 /// Called from GpuIndex for search
111 void searchImpl_(int n,
112 const float* x,
113 int k,
114 float* distances,
115 faiss::Index::idx_t* labels) const override;
116
117 private:
118 /// Checks user settings for consistency
119 void verifySettings_() const;
120
121 protected:
122 /// Our config object
123 const GpuIndexFlatConfig config_;
124
125 /// Holds our GPU data containing the list of vectors; is managed via raw
126 /// pointer so as to allow non-CUDA compilers to see this header
127 FlatIndex* data_;
128 };
129
130 /// Wrapper around the GPU implementation that looks like
131 /// faiss::IndexFlatL2; copies over centroid data from a given
132 /// faiss::IndexFlat
133 class GpuIndexFlatL2 : public GpuIndexFlat {
134 public:
135 /// Construct from a pre-existing faiss::IndexFlatL2 instance, copying
136 /// data over to the given GPU
137 GpuIndexFlatL2(GpuResources* resources,
138 faiss::IndexFlatL2* index,
139 GpuIndexFlatConfig config = GpuIndexFlatConfig());
140
141 /// Construct an empty instance that can be added to
142 GpuIndexFlatL2(GpuResources* resources,
143 int dims,
144 GpuIndexFlatConfig config = GpuIndexFlatConfig());
145
146 /// Initialize ourselves from the given CPU index; will overwrite
147 /// all data in ourselves
148 void copyFrom(faiss::IndexFlatL2* index);
149
150 /// Copy ourselves to the given CPU index; will overwrite all data
151 /// in the index instance
152 void copyTo(faiss::IndexFlatL2* index);
153 };
154
155 /// Wrapper around the GPU implementation that looks like
156 /// faiss::IndexFlatIP; copies over centroid data from a given
157 /// faiss::IndexFlat
158 class GpuIndexFlatIP : public GpuIndexFlat {
159 public:
160 /// Construct from a pre-existing faiss::IndexFlatIP instance, copying
161 /// data over to the given GPU
162 GpuIndexFlatIP(GpuResources* resources,
163 faiss::IndexFlatIP* index,
164 GpuIndexFlatConfig config = GpuIndexFlatConfig());
165
166 /// Construct an empty instance that can be added to
167 GpuIndexFlatIP(GpuResources* resources,
168 int dims,
169 GpuIndexFlatConfig config = GpuIndexFlatConfig());
170
171 /// Initialize ourselves from the given CPU index; will overwrite
172 /// all data in ourselves
173 void copyFrom(faiss::IndexFlatIP* index);
174
175 /// Copy ourselves to the given CPU index; will overwrite all data
176 /// in the index instance
177 void copyTo(faiss::IndexFlatIP* index);
178 };
179
180 } } // namespace
faiss::IndexFlat
Definition: IndexFlat.h:21
faiss::gpu::GpuIndexFlatConfig::useFloat16Accumulator
bool useFloat16Accumulator
Definition: GpuIndexFlat.h:38
faiss::gpu::GpuIndexFlatConfig::storeTransposed
bool storeTransposed
Definition: GpuIndexFlat.h:46
faiss::gpu::GpuIndexFlatL2::copyFrom
void copyFrom(faiss::IndexFlatL2 *index)
Definition: GpuIndexFlat.cu:297
faiss::gpu::GpuIndexFlat::copyTo
void copyTo(faiss::IndexFlat *index) const
Definition: GpuIndexFlat.cu:100
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:243
faiss::gpu::FlatIndex
Holder of GPU resources for a particular flat index.
Definition: FlatIndex.cuh:21
faiss::gpu::GpuIndexFlat::getGpuData
FlatIndex * getGpuData()
For internal access.
Definition: GpuIndexFlat.h:99
faiss::gpu::GpuIndexFlatL2::copyTo
void copyTo(faiss::IndexFlatL2 *index)
Definition: GpuIndexFlat.cu:302
faiss::IndexFlatL2
Definition: IndexFlat.h:77
faiss::gpu::GpuIndexFlat::getNumVecs
size_t getNumVecs() const
Returns the number of vectors we contain.
Definition: GpuIndexFlat.cu:124
faiss::gpu::GpuIndexFlat::searchImpl_
void searchImpl_(int n, const float *x, int k, float *distances, faiss::Index::idx_t *labels) const override
Called from GpuIndex for search.
Definition: GpuIndexFlat.cu:199
faiss::gpu::GpuResources
Definition: GpuResources.h:21
faiss::gpu::GpuIndexFlat::GpuIndexFlat
GpuIndexFlat(GpuResources *resources, const faiss::IndexFlat *index, GpuIndexFlatConfig config=GpuIndexFlatConfig())
Definition: GpuIndexFlat.cu:25
faiss::IndexFlatIP
Definition: IndexFlat.h:71
faiss::gpu::GpuIndexFlatConfig::useFloat16
bool useFloat16
Whether or not data is stored as float16.
Definition: GpuIndexFlat.h:33
faiss::gpu::GpuIndex
Definition: GpuIndex.h:33
faiss::gpu::GpuIndexFlatL2::GpuIndexFlatL2
GpuIndexFlatL2(GpuResources *resources, faiss::IndexFlatL2 *index, GpuIndexFlatConfig config=GpuIndexFlatConfig())
Definition: GpuIndexFlat.cu:284
faiss::Index::idx_t
long idx_t
all indices are this type
Definition: Index.h:62
faiss::gpu::GpuIndexFlatConfig
Definition: GpuIndexFlat.h:25
faiss::gpu::GpuIndexFlat::reconstruct
void reconstruct(faiss::Index::idx_t key, float *out) const override
Definition: GpuIndexFlat.cu:226
faiss::gpu::GpuIndexFlatIP::copyTo
void copyTo(faiss::IndexFlatIP *index)
Definition: GpuIndexFlat.cu:328
faiss::gpu::GpuIndexFlatIP::GpuIndexFlatIP
GpuIndexFlatIP(GpuResources *resources, faiss::IndexFlatIP *index, GpuIndexFlatConfig config=GpuIndexFlatConfig())
Definition: GpuIndexFlat.cu:310
faiss::gpu::GpuIndexFlat::addImpl_
void addImpl_(int n, const float *x, const Index::idx_t *ids) override
Called from GpuIndex for add.
Definition: GpuIndexFlat.cu:178
faiss::gpu::GpuIndexFlatIP
Definition: GpuIndexFlat.h:158
faiss::gpu::GpuIndexFlat::copyFrom
void copyFrom(const faiss::IndexFlat *index)
Definition: GpuIndexFlat.cu:67
faiss::gpu::GpuIndexFlat::addImplRequiresIDs_
bool addImplRequiresIDs_() const override
Flat index does not require IDs as there is no storage available for them.
Definition: GpuIndexFlat.cu:173
faiss::gpu::GpuIndexFlat::config_
const GpuIndexFlatConfig config_
Our config object.
Definition: GpuIndexFlat.h:123
faiss::gpu::GpuIndexFlatL2
Definition: GpuIndexFlat.h:133
faiss::gpu::GpuIndexConfig
Definition: GpuIndex.h:18
faiss::gpu::GpuIndexFlat::add
void add(faiss::Index::idx_t, const float *x) override
Overrides to avoid excessive copies.
Definition: GpuIndexFlat.cu:143
faiss::gpu::GpuIndexFlat
Definition: GpuIndexFlat.h:52
faiss::gpu::GpuIndexFlat::reset
void reset() override
Clears all vectors from this index.
Definition: GpuIndexFlat.cu:129
faiss::gpu::GpuIndexFlatIP::copyFrom
void copyFrom(faiss::IndexFlatIP *index)
Definition: GpuIndexFlat.cu:323
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:138
faiss::MetricType
MetricType
Some algorithms support both an inner product version and a L2 search version.
Definition: Index.h:44
faiss::gpu::GpuIndexFlat::data_
FlatIndex * data_
Definition: GpuIndexFlat.h:127