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 /// Wrapper around the GPU implementation that looks like
29 /// faiss::IndexFlat; copies over centroid data from a given
30 /// faiss::IndexFlat
31 class GpuIndexFlat : public GpuIndex {
32 public:
33 /// Construct from a pre-existing faiss::IndexFlat instance, copying
34 /// data over to the given GPU
35 GpuIndexFlat(GpuResources* resources,
36 int device,
37 bool useFloat16,
38 const faiss::IndexFlat* index);
39
40 /// Construct an empty instance that can be added to
41 GpuIndexFlat(GpuResources* resources,
42 int device,
43 int dims,
44 bool useFloat16,
45 faiss::MetricType metric);
46
47 ~GpuIndexFlat() override;
48
49 /// Set the minimum data size for searches (in MiB) for which we use
50 /// CPU -> GPU paging
51 void setMinPagingSize(size_t size);
52
53 /// Returns the current minimum data size for paged searches
54 size_t getMinPagingSize() const;
55
56 /// Do we store vectors and perform math in float16?
57 bool getUseFloat16() const;
58
59 /// Initialize ourselves from the given CPU index; will overwrite
60 /// all data in ourselves
61 void copyFrom(const faiss::IndexFlat* index);
62
63 /// Copy ourselves to the given CPU index; will overwrite all data
64 /// in the index instance
65 void copyTo(faiss::IndexFlat* index) const;
66
67 /// Returns the number of vectors we contain
68 size_t getNumVecs() const;
69
70 /// Clears all vectors from this index
71 void reset() override;
72
73 /// This index is not trained, so this does nothing
74 void train(Index::idx_t n, const float* x) override;
75
76 /// `x` can be resident on the CPU or any GPU; the proper copies are
77 /// performed
78 void add(Index::idx_t n, const float* x) override;
79
80 /// `x`, `distances` and `labels` can be resident on the CPU or any
81 /// GPU; copies are performed as needed
82 void search(faiss::Index::idx_t n,
83 const float* x,
84 faiss::Index::idx_t k,
85 float* distances,
86 faiss::Index::idx_t* labels) const 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(faiss::Index::idx_t i0,
94 faiss::Index::idx_t num,
95 float* out) const override;
96
97 void set_typename() override;
98
99 /// For internal access
100 inline FlatIndex* getGpuData() { return data_; }
101
102 protected:
103 /// Called from search when the input data is on the CPU;
104 /// potentially allows for pinned memory usage
105 void searchFromCpuPaged_(int n,
106 const float* x,
107 int k,
108 float* outDistancesData,
109 int* outIndicesData) const;
110
111 void searchNonPaged_(int n,
112 const float* x,
113 int k,
114 float* outDistancesData,
115 int* outIndicesData) const;
116
117 protected:
118 /// Size above which we page copies from the CPU to GPU
119 size_t minPagedSize_;
120
121 /// Whether or not we store our vectors in float32 or float16
122 const bool useFloat16_;
123
124 /// Holds our GPU data containing the list of vectors
125 FlatIndex* data_;
126 };
127
128 /// Wrapper around the GPU implementation that looks like
129 /// faiss::IndexFlatL2; copies over centroid data from a given
130 /// faiss::IndexFlat
131 class GpuIndexFlatL2 : public GpuIndexFlat {
132 public:
133 /// Construct from a pre-existing faiss::IndexFlatL2 instance, copying
134 /// data over to the given GPU
135 GpuIndexFlatL2(GpuResources* resources,
136 int device,
137 bool useFloat16,
138 faiss::IndexFlatL2* index);
139
140 /// Construct an empty instance that can be added to
141 GpuIndexFlatL2(GpuResources* resources,
142 int device,
143 int dims,
144 bool useFloat16);
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 int device,
164 bool useFloat16,
165 faiss::IndexFlatIP* index);
166
167 /// Construct an empty instance that can be added to
168 GpuIndexFlatIP(GpuResources* resources,
169 int device,
170 int dims,
171 bool useFloat16);
172
173 /// Initialize ourselves from the given CPU index; will overwrite
174 /// all data in ourselves
175 void copyFrom(faiss::IndexFlatIP* index);
176
177 /// Copy ourselves to the given CPU index; will overwrite all data
178 /// in the index instance
179 void copyTo(faiss::IndexFlatIP* index);
180 };
181
182 } } // namespace
faiss::IndexFlat
Definition: IndexFlat.h:24
faiss::gpu::GpuIndexFlat::GpuIndexFlat
GpuIndexFlat(GpuResources *resources, int device, bool useFloat16, const faiss::IndexFlat *index)
Definition: GpuIndexFlat.cu:33
faiss::gpu::GpuIndexFlatL2::copyFrom
void copyFrom(faiss::IndexFlatL2 *index)
Definition: GpuIndexFlat.cu:503
faiss::gpu::GpuIndexFlat::copyTo
void copyTo(faiss::IndexFlat *index) const
Definition: GpuIndexFlat.cu:107
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:457
faiss::gpu::GpuIndexFlat::getMinPagingSize
size_t getMinPagingSize() const
Returns the current minimum data size for paged searches.
Definition: GpuIndexFlat.cu:71
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:100
faiss::gpu::GpuIndexFlatL2::copyTo
void copyTo(faiss::IndexFlatL2 *index)
Definition: GpuIndexFlat.cu:508
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:268
faiss::gpu::GpuIndexFlat::useFloat16_
const bool useFloat16_
Whether or not we store our vectors in float32 or float16.
Definition: GpuIndexFlat.h:122
faiss::gpu::GpuIndexFlat::getNumVecs
size_t getNumVecs() const
Returns the number of vectors we contain.
Definition: GpuIndexFlat.cu:130
faiss::gpu::GpuResources
Definition: GpuResources.h:24
faiss::IndexFlatIP
Definition: IndexFlat.h:75
faiss::gpu::GpuIndexFlat::setMinPagingSize
void setMinPagingSize(size_t size)
Definition: GpuIndexFlat.cu:66
faiss::gpu::GpuIndex
Definition: GpuIndex.h:20
faiss::gpu::GpuIndexFlat::getUseFloat16
bool getUseFloat16() const
Do we store vectors and perform math in float16?
Definition: GpuIndexFlat.cu:76
faiss::gpu::GpuIndexFlatL2::GpuIndexFlatL2
GpuIndexFlatL2(GpuResources *resources, int device, bool useFloat16, faiss::IndexFlatL2 *index)
Definition: GpuIndexFlat.cu:488
faiss::gpu::GpuIndexFlat::reconstruct
void reconstruct(faiss::Index::idx_t key, float *out) const override
Definition: GpuIndexFlat.cu:440
faiss::gpu::GpuIndexFlatIP::copyTo
void copyTo(faiss::IndexFlatIP *index)
Definition: GpuIndexFlat.cu:537
faiss::Index::idx_t
long idx_t
all indices are this type
Definition: Index.h:64
faiss::gpu::GpuIndexFlatIP
Definition: GpuIndexFlat.h:158
faiss::gpu::GpuIndexFlat::copyFrom
void copyFrom(const faiss::IndexFlat *index)
Definition: GpuIndexFlat.cu:81
faiss::gpu::GpuIndexFlat::add
void add(Index::idx_t n, const float *x) override
Definition: GpuIndexFlat.cu:149
faiss::gpu::GpuIndexFlat::minPagedSize_
size_t minPagedSize_
Size above which we page copies from the CPU to GPU.
Definition: GpuIndexFlat.h:119
faiss::gpu::GpuIndexFlatIP::GpuIndexFlatIP
GpuIndexFlatIP(GpuResources *resources, int device, bool useFloat16, faiss::IndexFlatIP *index)
Definition: GpuIndexFlat.cu:516
faiss::gpu::GpuIndexFlatL2
Definition: GpuIndexFlat.h:131
faiss::gpu::GpuIndexFlat
Definition: GpuIndexFlat.h:31
faiss::gpu::GpuIndexFlat::reset
void reset() override
Clears all vectors from this index.
Definition: GpuIndexFlat.cu:135
faiss::gpu::GpuIndexFlatIP::copyFrom
void copyFrom(faiss::IndexFlatIP *index)
Definition: GpuIndexFlat.cu:532
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:167
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:144
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:125