Faiss
 All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
Distance.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 "Distance.cuh"
13 #include "BroadcastSum.cuh"
14 #include "L2Norm.cuh"
15 #include "L2Select.cuh"
16 #include "../../FaissAssert.h"
17 #include "../GpuResources.h"
18 #include "../utils/DeviceUtils.h"
19 #include "../utils/Limits.cuh"
20 #include "../utils/MatrixMult.cuh"
21 #include "../utils/BlockSelectKernel.cuh"
22 
23 #include <memory>
24 #include <thrust/fill.h>
25 #include <thrust/for_each.h>
26 #include <thrust/device_ptr.h>
27 #include <thrust/execution_policy.h>
28 
29 namespace faiss { namespace gpu {
30 
31 constexpr int kDefaultTileSize = 256;
32 
33 template <typename T>
34 void runL2Distance(GpuResources* resources,
35  Tensor<T, 2, true>& centroids,
36  Tensor<T, 1, true>* centroidNorms,
37  Tensor<T, 2, true>& queries,
38  int k,
39  Tensor<T, 2, true>& outDistances,
40  Tensor<int, 2, true>& outIndices,
41  bool ignoreOutDistances = false,
42  int tileSize = -1) {
43  FAISS_ASSERT(outDistances.getSize(0) == queries.getSize(0));
44  FAISS_ASSERT(outIndices.getSize(0) == queries.getSize(0));
45  FAISS_ASSERT(outDistances.getSize(1) == k);
46  FAISS_ASSERT(outIndices.getSize(1) == k);
47 
48  auto& mem = resources->getMemoryManagerCurrentDevice();
49  auto defaultStream = resources->getDefaultStreamCurrentDevice();
50 
51  // If we're quering against a 0 sized set, just return empty results
52  if (centroids.numElements() == 0) {
53  thrust::fill(thrust::cuda::par.on(defaultStream),
54  outDistances.data(), outDistances.end(),
55  Limits<T>::getMax());
56 
57  thrust::fill(thrust::cuda::par.on(defaultStream),
58  outIndices.data(), outIndices.end(),
59  -1);
60 
61  return;
62  }
63 
64  // If ||c||^2 is not pre-computed, calculate it
65  DeviceTensor<T, 1, true> cNorms;
66  if (!centroidNorms) {
67  cNorms = std::move(DeviceTensor<T, 1, true>(
68  mem,
69  {centroids.getSize(0)}, defaultStream));
70  runL2Norm(centroids, cNorms, true, defaultStream);
71  centroidNorms = &cNorms;
72  }
73 
74  //
75  // Prepare norm vector ||q||^2; ||c||^2 is already pre-computed
76  //
77  int qNormSize[1] = {queries.getSize(0)};
78  DeviceTensor<T, 1, true> queryNorms(mem, qNormSize, defaultStream);
79 
80  // ||q||^2
81  runL2Norm(queries, queryNorms, true, defaultStream);
82 
83  //
84  // Handle the problem in row tiles, to avoid excessive temporary
85  // memory requests
86  //
87 
88  FAISS_ASSERT(k <= centroids.getSize(0));
89  FAISS_ASSERT(k <= 1024); // select limitation
90 
91  // To allocate all of (#queries, #centroids) is potentially too much
92  // memory. Limit our total size requested
93  size_t distanceRowSize = centroids.getSize(0) * sizeof(T);
94 
95  // FIXME: parameterize based on # of centroids and DeviceMemory
96  int defaultTileSize = sizeof(T) < 4 ? kDefaultTileSize * 2 : kDefaultTileSize;
97  tileSize = tileSize <= 0 ? defaultTileSize : tileSize;
98 
99  int maxQueriesPerIteration = std::min(tileSize, queries.getSize(0));
100 
101  // Temporary output memory space we'll use
102  DeviceTensor<T, 2, true> distanceBuf1(
103  mem, {maxQueriesPerIteration, centroids.getSize(0)}, defaultStream);
104  DeviceTensor<T, 2, true> distanceBuf2(
105  mem, {maxQueriesPerIteration, centroids.getSize(0)}, defaultStream);
106  DeviceTensor<T, 2, true>* distanceBufs[2] =
107  {&distanceBuf1, &distanceBuf2};
108 
109  auto streams = resources->getAlternateStreamsCurrentDevice();
110  streamWait(streams, {defaultStream});
111 
112  int curStream = 0;
113 
114  for (int i = 0; i < queries.getSize(0); i += maxQueriesPerIteration) {
115  int numQueriesForIteration = std::min(maxQueriesPerIteration,
116  queries.getSize(0) - i);
117 
118  auto distanceBufView =
119  distanceBufs[curStream]->narrowOutermost(0, numQueriesForIteration);
120  auto queryView =
121  queries.narrowOutermost(i, numQueriesForIteration);
122  auto outDistanceView =
123  outDistances.narrowOutermost(i, numQueriesForIteration);
124  auto outIndexView =
125  outIndices.narrowOutermost(i, numQueriesForIteration);
126  auto queryNormNiew =
127  queryNorms.narrowOutermost(i, numQueriesForIteration);
128 
129  // L2 distance is ||c||^2 - 2qc + ||q||^2
130 
131  // -2qc
132  // (query id x dim) x (centroid id, dim)' = (query id, centroid id)
133  runMatrixMult(distanceBufView, false,
134  queryView, false,
135  centroids, true,
136  -2.0f, 0.0f,
137  resources->getBlasHandleCurrentDevice(),
138  streams[curStream]);
139 
140  // For L2 distance, we use this fused kernel that performs both
141  // adding ||c||^2 to -2qc and k-selection, so we only need two
142  // passes (one write by the gemm, one read here) over the huge
143  // region of output memory
144  runL2SelectMin(distanceBufView,
145  *centroidNorms,
146  outDistanceView,
147  outIndexView,
148  k,
149  streams[curStream]);
150 
151  if (!ignoreOutDistances) {
152  // expand (query id) to (query id, k) by duplicating along rows
153  // top-k ||c||^2 - 2qc + ||q||^2 in the form (query id, k)
154  runSumAlongRows(queryNormNiew, outDistanceView, streams[curStream]);
155  }
156 
157  curStream = (curStream + 1) % 2;
158  }
159 
160  // Have the desired ordering stream wait on the multi-stream
161  streamWait({defaultStream}, streams);
162 }
163 
164 template <typename T>
165 void runIPDistance(GpuResources* resources,
166  Tensor<T, 2, true>& centroids,
167  Tensor<T, 2, true>& queries,
168  int k,
169  Tensor<T, 2, true>& outDistances,
170  Tensor<int, 2, true>& outIndices,
171  int tileSize = -1) {
172  FAISS_ASSERT(outDistances.getSize(0) == queries.getSize(0));
173  FAISS_ASSERT(outIndices.getSize(0) == queries.getSize(0));
174  FAISS_ASSERT(outDistances.getSize(1) == k);
175  FAISS_ASSERT(outIndices.getSize(1) == k);
176 
177  auto& mem = resources->getMemoryManagerCurrentDevice();
178  auto defaultStream = resources->getDefaultStreamCurrentDevice();
179 
180  // If we're quering against a 0 sized set, just return empty results
181  if (centroids.numElements() == 0) {
182  thrust::fill(thrust::cuda::par.on(defaultStream),
183  outDistances.data(), outDistances.end(),
184  Limits<T>::getMax());
185 
186  thrust::fill(thrust::cuda::par.on(defaultStream),
187  outIndices.data(), outIndices.end(),
188  -1);
189 
190  return;
191  }
192 
193  //
194  // Handle the problem in row tiles, to avoid excessive temporary
195  // memory requests
196  //
197 
198  FAISS_ASSERT(k <= centroids.getSize(0));
199  FAISS_ASSERT(k <= 1024); // select limitation
200 
201  // To allocate all of (#queries, #centroids) is potentially too much
202  // memory. Limit our total size requested
203  size_t distanceRowSize = centroids.getSize(0) * sizeof(T);
204 
205  // FIXME: parameterize based on # of centroids and DeviceMemory
206  int defaultTileSize = sizeof(T) < 4 ? kDefaultTileSize * 2 : kDefaultTileSize;
207  tileSize = tileSize <= 0 ? defaultTileSize : tileSize;
208 
209  int maxQueriesPerIteration = std::min(tileSize, queries.getSize(0));
210 
211  // Temporary output memory space we'll use
212  DeviceTensor<T, 2, true> distanceBuf1(
213  mem, {maxQueriesPerIteration, centroids.getSize(0)}, defaultStream);
214  DeviceTensor<T, 2, true> distanceBuf2(
215  mem, {maxQueriesPerIteration, centroids.getSize(0)}, defaultStream);
216  DeviceTensor<T, 2, true>* distanceBufs[2] =
217  {&distanceBuf1, &distanceBuf2};
218 
219  auto streams = resources->getAlternateStreamsCurrentDevice();
220  streamWait(streams, {defaultStream});
221 
222  int curStream = 0;
223 
224  for (int i = 0; i < queries.getSize(0); i += maxQueriesPerIteration) {
225  int numQueriesForIteration = std::min(maxQueriesPerIteration,
226  queries.getSize(0) - i);
227 
228  auto distanceBufView =
229  distanceBufs[curStream]->narrowOutermost(0, numQueriesForIteration);
230  auto queryView =
231  queries.narrowOutermost(i, numQueriesForIteration);
232  auto outDistanceView =
233  outDistances.narrowOutermost(i, numQueriesForIteration);
234  auto outIndexView =
235  outIndices.narrowOutermost(i, numQueriesForIteration);
236 
237  // (query id x dim) x (centroid id, dim)' = (query id, centroid id)
238  runMatrixMult(distanceBufView, false,
239  queryView, false, centroids, true,
240  1.0f, 0.0f,
241  resources->getBlasHandleCurrentDevice(),
242  streams[curStream]);
243 
244  // top-k of dot products
245  // (query id, top k centroids)
246  runBlockSelect(distanceBufView,
247  outDistanceView,
248  outIndexView,
249  true, k, streams[curStream]);
250 
251  curStream = (curStream + 1) % 2;
252  }
253 
254  streamWait({defaultStream}, streams);
255 }
256 
257 //
258 // Instantiations of the distance templates
259 //
260 
261 void
262 runIPDistance(GpuResources* resources,
263  Tensor<float, 2, true>& vectors,
264  Tensor<float, 2, true>& queries,
265  int k,
266  Tensor<float, 2, true>& outDistances,
267  Tensor<int, 2, true>& outIndices,
268  int tileSize) {
269  runIPDistance<float>(resources,
270  vectors,
271  queries,
272  k,
273  outDistances,
274  outIndices,
275  tileSize);
276 }
277 
278 #ifdef FAISS_USE_FLOAT16
279 void
280 runIPDistance(GpuResources* resources,
281  Tensor<half, 2, true>& vectors,
282  Tensor<half, 2, true>& queries,
283  int k,
284  Tensor<half, 2, true>& outDistances,
285  Tensor<int, 2, true>& outIndices,
286  int tileSize) {
287  runIPDistance<half>(resources,
288  vectors,
289  queries,
290  k,
291  outDistances,
292  outIndices,
293  tileSize);
294 }
295 #endif
296 
297 void
298 runL2Distance(GpuResources* resources,
299  Tensor<float, 2, true>& vectors,
300  Tensor<float, 1, true>* vectorNorms,
301  Tensor<float, 2, true>& queries,
302  int k,
303  Tensor<float, 2, true>& outDistances,
304  Tensor<int, 2, true>& outIndices,
305  bool ignoreOutDistances,
306  int tileSize) {
307  runL2Distance<float>(resources,
308  vectors,
309  vectorNorms,
310  queries,
311  k,
312  outDistances,
313  outIndices,
314  ignoreOutDistances,
315  tileSize);
316 }
317 
318 #ifdef FAISS_USE_FLOAT16
319 void
320 runL2Distance(GpuResources* resources,
321  Tensor<half, 2, true>& vectors,
322  Tensor<half, 1, true>* vectorNorms,
323  Tensor<half, 2, true>& queries,
324  int k,
325  Tensor<half, 2, true>& outDistances,
326  Tensor<int, 2, true>& outIndices,
327  bool ignoreOutDistances,
328  int tileSize) {
329  runL2Distance<half>(resources,
330  vectors,
331  vectorNorms,
332  queries,
333  k,
334  outDistances,
335  outIndices,
336  ignoreOutDistances,
337  tileSize);
338 }
339 #endif
340 
341 } } // namespace
Generated by   doxygen 1.8.5