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