Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
FlatIndex.cu
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 #include "FlatIndex.cuh"
13 #include "Distance.cuh"
14 #include "L2Norm.cuh"
15 #include "../utils/CopyUtils.cuh"
16 #include "../utils/DeviceUtils.h"
17
18 namespace faiss { namespace gpu {
19
20 FlatIndex::FlatIndex(GpuResources* res,
21 int dim,
22 bool l2Distance,
23 bool useFloat16) :
24 resources_(res),
25 dim_(dim),
26 useFloat16_(useFloat16),
27 l2Distance_(l2Distance),
28 num_(0) {
29 #ifndef FAISS_USE_FLOAT16
30 FAISS_ASSERT(!useFloat16_);
31 #endif
32 }
33
34 bool
35 FlatIndex::getUseFloat16() const {
36 return useFloat16_;
37 }
38
39 /// Returns the number of vectors we contain
40 int FlatIndex::getSize() const {
41 #ifdef FAISS_USE_FLOAT16
42 if (useFloat16_) {
43 return vectorsHalf_.getSize(0);
44 }
45 #endif
46
47 return vectors_.getSize(0);
48 }
49
50 int FlatIndex::getDim() const {
51 #ifdef FAISS_USE_FLOAT16
52 if (useFloat16_) {
53 return vectorsHalf_.getSize(1);
54 }
55 #endif
56
57 return vectors_.getSize(1);
58 }
59
60 Tensor<float, 2, true>&
61 FlatIndex::getVectorsFloat32Ref() {
62 return vectors_;
63 }
64
65 #ifdef FAISS_USE_FLOAT16
66 Tensor<half, 2, true>&
67 FlatIndex::getVectorsFloat16Ref() {
68 return vectorsHalf_;
69 }
70 #endif
71
72 DeviceTensor<float, 2, true>
73 FlatIndex::getVectorsFloat32Copy(cudaStream_t stream) {
74 return getVectorsFloat32Copy(0, num_, stream);
75 }
76
77 DeviceTensor<float, 2, true>
78 FlatIndex::getVectorsFloat32Copy(int from, int num, cudaStream_t stream) {
79 DeviceTensor<float, 2, true> vecFloat32({num, dim_});
80
81 if (useFloat16_) {
82 #ifdef FAISS_USE_FLOAT16
83 runConvertToFloat32(vecFloat32.data(),
84 vectorsHalf_[from].data(),
85 num * dim_, stream);
86 #endif
87 } else {
88 vectors_.copyTo(vecFloat32, stream);
89 }
90
91 return vecFloat32;
92 }
93
94 void
95 FlatIndex::query(Tensor<float, 2, true>& vecs,
96 int k,
97 Tensor<float, 2, true>& outDistances,
98 Tensor<int, 2, true>& outIndices,
99 bool exactDistance,
100 int tileSize) {
101 auto stream = resources_->getDefaultStreamCurrentDevice();
102 auto& mem = resources_->getMemoryManagerCurrentDevice();
103
104 if (useFloat16_) {
105 // We need to convert to float16
106 #ifdef FAISS_USE_FLOAT16
107 auto vecsHalf = toHalf<2>(resources_, stream, vecs);
108
109 DeviceTensor<half, 2, true> outDistancesHalf(
110 mem, {outDistances.getSize(0), outDistances.getSize(1)}, stream);
111
112 query(vecsHalf, k, outDistancesHalf, outIndices, exactDistance, tileSize);
113
114 if (exactDistance) {
115 // Convert outDistances back
116 fromHalf<2>(stream, outDistancesHalf, outDistances);
117 }
118 #endif
119 } else {
120 if (l2Distance_) {
121 runL2Distance(resources_,
122 vectors_,
123 &norms_,
124 vecs,
125 k,
126 outDistances,
127 outIndices,
128 // FIXME
129 !exactDistance,
130 tileSize);
131 } else {
132 runIPDistance(resources_,
133 vectors_,
134 vecs,
135 k,
136 outDistances,
137 outIndices,
138 tileSize);
139 }
140 }
141 }
142
143 #ifdef FAISS_USE_FLOAT16
144 void
145 FlatIndex::query(Tensor<half, 2, true>& vecs,
146 int k,
147 Tensor<half, 2, true>& outDistances,
148 Tensor<int, 2, true>& outIndices,
149 bool exactDistance,
150 int tileSize) {
151 FAISS_ASSERT(useFloat16_);
152
153 if (l2Distance_) {
154 runL2Distance(resources_,
155 vectorsHalf_,
156 &normsHalf_,
157 vecs,
158 k,
159 outDistances,
160 outIndices,
161 // FIXME
162 !exactDistance,
163 tileSize);
164 } else {
165 runIPDistance(resources_,
166 vectorsHalf_,
167 vecs,
168 k,
169 outDistances,
170 outIndices,
171 tileSize);
172 }
173 }
174 #endif
175
176 void
177 FlatIndex::add(const float* data, int numVecs, cudaStream_t stream) {
178 if (numVecs == 0) {
179 return;
180 }
181
182 if (useFloat16_) {
183 #ifdef FAISS_USE_FLOAT16
184 // Make sure that `data` is on our device; we'll run the
185 // conversion on our device
186 auto devData = toDevice<float, 2>(resources_,
187 getCurrentDevice(),
188 (float*) data,
189 stream,
190 {numVecs, dim_});
191
192 auto devDataHalf = toHalf<2>(resources_, stream, devData);
193
194 rawData_.append((char*) devDataHalf.data(),
195 devDataHalf.getSizeInBytes(),
196 stream);
197 #endif
198 } else {
199 rawData_.append((char*) data,
200 (size_t) dim_ * numVecs * sizeof(float),
201 stream);
202 }
203
204 num_ += numVecs;
205
206 if (useFloat16_) {
207 #ifdef FAISS_USE_FLOAT16
208 DeviceTensor<half, 2, true> vectorsHalf(
209 (half*) rawData_.data(), {(int) num_, dim_});
210 vectorsHalf_ = std::move(vectorsHalf);
211 #endif
212 } else {
213 DeviceTensor<float, 2, true> vectors(
214 (float*) rawData_.data(), {(int) num_, dim_});
215 vectors_ = std::move(vectors);
216 }
217
218 if (l2Distance_) {
219 // Precompute L2 norms of our database
220 if (useFloat16_) {
221 #ifdef FAISS_USE_FLOAT16
222 DeviceTensor<half, 1, true> normsHalf({(int) num_});
223 runL2Norm(vectorsHalf_, normsHalf, true, stream);
224 normsHalf_ = std::move(normsHalf);
225 #endif
226 } else {
227 DeviceTensor<float, 1, true> norms({(int) num_});
228 runL2Norm(vectors_, norms, true, stream);
229 norms_ = std::move(norms);
230 }
231 }
232 }
233
234 void
235 FlatIndex::reset() {
236 rawData_.clear();
237 vectors_ = std::move(DeviceTensor<float, 2, true>());
238 norms_ = std::move(DeviceTensor<float, 1, true>());
239 num_ = 0;
240 }
241
242 } }
faiss::gpu::FlatIndex::getVectorsFloat32Copy
DeviceTensor< float, 2, true > getVectorsFloat32Copy(cudaStream_t stream)
Definition: FlatIndex.cu:73
faiss::gpu::FlatIndex::getSize
int getSize() const
Returns the number of vectors we contain.
Definition: FlatIndex.cu:40
faiss::gpu::Tensor::copyTo
__host__ void copyTo(Tensor< T, Dim, Contig, IndexT, PtrTraits > &t, cudaStream_t stream)
Copies ourselves into a tensor; sizes must match.
Definition: Tensor-inl.cuh:139
faiss::gpu::FlatIndex::add
void add(const float *data, int numVecs, cudaStream_t stream)
Definition: FlatIndex.cu:177
faiss::gpu::Tensor
Our tensor type.
Definition: Tensor.cuh:31
faiss::gpu::Tensor::getSize
__host__ __device__ IndexT getSize(int i) const
Definition: Tensor.cuh:210
faiss::gpu::FlatIndex::getVectorsFloat32Ref
Tensor< float, 2, true > & getVectorsFloat32Ref()
Returns a reference to our vectors currently in use.
Definition: FlatIndex.cu:61
faiss::gpu::FlatIndex::reset
void reset()
Free all storage.
Definition: FlatIndex.cu:235
faiss::gpu::DeviceTensor< float, 2, true >