Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
GpuIndexFlat.h
1
2 /**
3 * Copyright (c) 2015-present, Facebook, Inc.
4 * All rights reserved.
5 *
6 * This source code is licensed under the CC-by-NC license found in the
7 * LICENSE file in the root directory of this source tree.
8 */
9
10 // Copyright 2004-present Facebook. All Rights Reserved.
11
12 #pragma once
13
14 #include "GpuIndex.h"
15
16 namespace faiss {
17
18 struct IndexFlat;
19 struct IndexFlatL2;
20 struct IndexFlatIP;
21
22 }
23
24 namespace faiss { namespace gpu {
25
26 struct FlatIndex;
27
28 struct GpuIndexFlatConfig {
29 inline GpuIndexFlatConfig()
30 : device(0),
31 useFloat16(false),
32 storeTransposed(false) {
33 }
34
35 int device;
36 bool useFloat16;
37 bool storeTransposed;
38 };
39
40 /// Wrapper around the GPU implementation that looks like
41 /// faiss::IndexFlat; copies over centroid data from a given
42 /// faiss::IndexFlat
43 class GpuIndexFlat : public GpuIndex {
44 public:
45 /// Construct from a pre-existing faiss::IndexFlat instance, copying
46 /// data over to the given GPU
47 GpuIndexFlat(GpuResources* resources,
48 const faiss::IndexFlat* index,
49 GpuIndexFlatConfig config = GpuIndexFlatConfig());
50
51 /// Construct an empty instance that can be added to
52 GpuIndexFlat(GpuResources* resources,
53 int dims,
54 faiss::MetricType metric,
55 GpuIndexFlatConfig config = GpuIndexFlatConfig());
56
57 virtual ~GpuIndexFlat();
58
59 /// Set the minimum data size for searches (in MiB) for which we use
60 /// CPU -> GPU paging
61 void setMinPagingSize(size_t size);
62
63 /// Returns the current minimum data size for paged searches
64 size_t getMinPagingSize() const;
65
66 /// Do we store vectors and perform math in float16?
67 bool getUseFloat16() const;
68
69 /// Initialize ourselves from the given CPU index; will overwrite
70 /// all data in ourselves
71 void copyFrom(const faiss::IndexFlat* index);
72
73 /// Copy ourselves to the given CPU index; will overwrite all data
74 /// in the index instance
75 void copyTo(faiss::IndexFlat* index) const;
76
77 /// Returns the number of vectors we contain
78 size_t getNumVecs() const;
79
80 /// Clears all vectors from this index
81 virtual void reset();
82
83 /// This index is not trained, so this does nothing
84 virtual void train(Index::idx_t n, const float* x);
85
86 /// `x`, `distances` and `labels` can be resident on the CPU or any
87 /// GPU; copies are performed as needed
88 /// We have our own implementation here which handles CPU async
89 /// copies; searchImpl_ is not called
90 /// FIXME: move paged impl into GpuIndex
91 virtual void search(faiss::Index::idx_t n,
92 const float* x,
93 faiss::Index::idx_t k,
94 float* distances,
95 faiss::Index::idx_t* labels) const;
96
97 /// Reconstruction methods; prefer the batch reconstruct as it will
98 /// be more efficient
99 virtual void reconstruct(faiss::Index::idx_t key, float* out) const;
100
101 /// Batch reconstruction method
102 virtual void reconstruct_n(faiss::Index::idx_t i0,
103 faiss::Index::idx_t num,
104 float* out) const;
105
106 virtual void set_typename();
107
108 /// For internal access
109 inline FlatIndex* getGpuData() { return data_; }
110
111 protected:
112 /// Called from GpuIndex for add
113 virtual void addImpl_(faiss::Index::idx_t n,
114 const float* x,
115 const faiss::Index::idx_t* ids);
116
117 /// Should not be called (we have our own implementation)
118 virtual void searchImpl_(faiss::Index::idx_t n,
119 const float* x,
120 faiss::Index::idx_t k,
121 float* distances,
122 faiss::Index::idx_t* labels) const;
123
124 /// Called from search when the input data is on the CPU;
125 /// potentially allows for pinned memory usage
126 void searchFromCpuPaged_(int n,
127 const float* x,
128 int k,
129 float* outDistancesData,
130 int* outIndicesData) const;
131
132 void searchNonPaged_(int n,
133 const float* x,
134 int k,
135 float* outDistancesData,
136 int* outIndicesData) const;
137
138 protected:
139 /// Size above which we page copies from the CPU to GPU
140 size_t minPagedSize_;
141
142 const GpuIndexFlatConfig config_;
143
144 /// Holds our GPU data containing the list of vectors
145 FlatIndex* data_;
146 };
147
148 /// Wrapper around the GPU implementation that looks like
149 /// faiss::IndexFlatL2; copies over centroid data from a given
150 /// faiss::IndexFlat
151 class GpuIndexFlatL2 : public GpuIndexFlat {
152 public:
153 /// Construct from a pre-existing faiss::IndexFlatL2 instance, copying
154 /// data over to the given GPU
155 GpuIndexFlatL2(GpuResources* resources,
156 faiss::IndexFlatL2* index,
157 GpuIndexFlatConfig config = GpuIndexFlatConfig());
158
159 /// Construct an empty instance that can be added to
160 GpuIndexFlatL2(GpuResources* resources,
161 int dims,
162 GpuIndexFlatConfig config = GpuIndexFlatConfig());
163
164 /// Initialize ourselves from the given CPU index; will overwrite
165 /// all data in ourselves
166 void copyFrom(faiss::IndexFlatL2* index);
167
168 /// Copy ourselves to the given CPU index; will overwrite all data
169 /// in the index instance
170 void copyTo(faiss::IndexFlatL2* index);
171 };
172
173 /// Wrapper around the GPU implementation that looks like
174 /// faiss::IndexFlatIP; copies over centroid data from a given
175 /// faiss::IndexFlat
176 class GpuIndexFlatIP : public GpuIndexFlat {
177 public:
178 /// Construct from a pre-existing faiss::IndexFlatIP instance, copying
179 /// data over to the given GPU
180 GpuIndexFlatIP(GpuResources* resources,
181 faiss::IndexFlatIP* index,
182 GpuIndexFlatConfig config = GpuIndexFlatConfig());
183
184 /// Construct an empty instance that can be added to
185 GpuIndexFlatIP(GpuResources* resources,
186 int dims,
187 GpuIndexFlatConfig config = GpuIndexFlatConfig());
188
189 /// Initialize ourselves from the given CPU index; will overwrite
190 /// all data in ourselves
191 void copyFrom(faiss::IndexFlatIP* index);
192
193 /// Copy ourselves to the given CPU index; will overwrite all data
194 /// in the index instance
195 void copyTo(faiss::IndexFlatIP* index);
196 };
197
198 } } // 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