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 namespace {
32 
33 constexpr int kDefaultTileSize = 256;
34 
35 template <typename T>
36 int chooseTileSize(int tileSizeOverride,
37  size_t numCentroids,
38  size_t tempMemAvailable) {
39  if (tileSizeOverride > 0) {
40  return tileSizeOverride;
41  }
42 
43  size_t tileSize =
44  sizeof(T) < 4 ? kDefaultTileSize * 2 : kDefaultTileSize;
45 
46  while (tileSize > 64) {
47  size_t memRequirement = 2 * tileSize * numCentroids * sizeof(T);
48 
49  if (memRequirement <= tempMemAvailable) {
50  // This fits entirely into our temporary memory
51  return tileSize;
52  }
53 
54  // Otherwise, halve the tile size
55  tileSize /= 2;
56  }
57 
58  // We use 64 as the minimum acceptable tile size
59  FAISS_ASSERT(tileSize >= 64);
60 
61  // FIXME: if we're running with no available temp memory, do we try
62  // and go larger based on free memory available on the device?
63 
64  return tileSize;
65 }
66 
67 }
68 
69 template <typename T>
70 void runL2Distance(GpuResources* resources,
71  Tensor<T, 2, true>& centroids,
72  Tensor<T, 2, true>* centroidsTransposed,
73  Tensor<T, 1, true>* centroidNorms,
74  Tensor<T, 2, true>& queries,
75  int k,
76  Tensor<T, 2, true>& outDistances,
77  Tensor<int, 2, true>& outIndices,
78  bool ignoreOutDistances = false,
79  int tileSizeOverride = -1) {
80  FAISS_ASSERT(outDistances.getSize(0) == queries.getSize(0));
81  FAISS_ASSERT(outIndices.getSize(0) == queries.getSize(0));
82  FAISS_ASSERT(outDistances.getSize(1) == k);
83  FAISS_ASSERT(outIndices.getSize(1) == k);
84 
85  auto& mem = resources->getMemoryManagerCurrentDevice();
86  auto defaultStream = resources->getDefaultStreamCurrentDevice();
87 
88  // If we're quering against a 0 sized set, just return empty results
89  if (centroids.numElements() == 0) {
90  thrust::fill(thrust::cuda::par.on(defaultStream),
91  outDistances.data(), outDistances.end(),
92  Limits<T>::getMax());
93 
94  thrust::fill(thrust::cuda::par.on(defaultStream),
95  outIndices.data(), outIndices.end(),
96  -1);
97 
98  return;
99  }
100 
101  // If ||c||^2 is not pre-computed, calculate it
102  DeviceTensor<T, 1, true> cNorms;
103  if (!centroidNorms) {
104  cNorms = std::move(DeviceTensor<T, 1, true>(
105  mem,
106  {centroids.getSize(0)}, defaultStream));
107  runL2Norm(centroids, cNorms, true, defaultStream);
108  centroidNorms = &cNorms;
109  }
110 
111  //
112  // Prepare norm vector ||q||^2; ||c||^2 is already pre-computed
113  //
114  int qNormSize[1] = {queries.getSize(0)};
115  DeviceTensor<T, 1, true> queryNorms(mem, qNormSize, defaultStream);
116 
117  // ||q||^2
118  runL2Norm(queries, queryNorms, true, defaultStream);
119 
120  //
121  // Handle the problem in row tiles, to avoid excessive temporary
122  // memory requests
123  //
124 
125  FAISS_ASSERT(k <= centroids.getSize(0));
126  FAISS_ASSERT(k <= 1024); // select limitation
127 
128  int tileSize =
129  chooseTileSize<T>(
130  tileSizeOverride,
131  centroids.getSize(0),
132  resources->getMemoryManagerCurrentDevice().getSizeAvailable());
133 
134  int maxQueriesPerIteration = std::min(tileSize, queries.getSize(0));
135 
136  // Temporary output memory space we'll use
137  DeviceTensor<T, 2, true> distanceBuf1(
138  mem, {maxQueriesPerIteration, centroids.getSize(0)}, defaultStream);
139  DeviceTensor<T, 2, true> distanceBuf2(
140  mem, {maxQueriesPerIteration, centroids.getSize(0)}, defaultStream);
141  DeviceTensor<T, 2, true>* distanceBufs[2] =
142  {&distanceBuf1, &distanceBuf2};
143 
144  auto streams = resources->getAlternateStreamsCurrentDevice();
145  streamWait(streams, {defaultStream});
146 
147  int curStream = 0;
148 
149  for (int i = 0; i < queries.getSize(0); i += maxQueriesPerIteration) {
150  int numQueriesForIteration = std::min(maxQueriesPerIteration,
151  queries.getSize(0) - i);
152 
153  auto distanceBufView =
154  distanceBufs[curStream]->narrowOutermost(0, numQueriesForIteration);
155  auto queryView =
156  queries.narrowOutermost(i, numQueriesForIteration);
157  auto outDistanceView =
158  outDistances.narrowOutermost(i, numQueriesForIteration);
159  auto outIndexView =
160  outIndices.narrowOutermost(i, numQueriesForIteration);
161  auto queryNormNiew =
162  queryNorms.narrowOutermost(i, numQueriesForIteration);
163 
164  // L2 distance is ||c||^2 - 2qc + ||q||^2
165 
166  // -2qc
167  // (query id x dim) x (centroid id, dim)' = (query id, centroid id)
168  runMatrixMult(distanceBufView, false,
169  queryView, false,
170  centroidsTransposed ? *centroidsTransposed : centroids,
171  centroidsTransposed ? false : true,
172  -2.0f, 0.0f,
173  resources->getBlasHandleCurrentDevice(),
174  streams[curStream]);
175 
176  // For L2 distance, we use this fused kernel that performs both
177  // adding ||c||^2 to -2qc and k-selection, so we only need two
178  // passes (one write by the gemm, one read here) over the huge
179  // region of output memory
180  runL2SelectMin(distanceBufView,
181  *centroidNorms,
182  outDistanceView,
183  outIndexView,
184  k,
185  streams[curStream]);
186 
187  if (!ignoreOutDistances) {
188  // expand (query id) to (query id, k) by duplicating along rows
189  // top-k ||c||^2 - 2qc + ||q||^2 in the form (query id, k)
190  runSumAlongRows(queryNormNiew, outDistanceView, streams[curStream]);
191  }
192 
193  curStream = (curStream + 1) % 2;
194  }
195 
196  // Have the desired ordering stream wait on the multi-stream
197  streamWait({defaultStream}, streams);
198 }
199 
200 template <typename T>
201 void runIPDistance(GpuResources* resources,
202  Tensor<T, 2, true>& centroids,
203  Tensor<T, 2, true>* centroidsTransposed,
204  Tensor<T, 2, true>& queries,
205  int k,
206  Tensor<T, 2, true>& outDistances,
207  Tensor<int, 2, true>& outIndices,
208  int tileSizeOverride = -1) {
209  FAISS_ASSERT(outDistances.getSize(0) == queries.getSize(0));
210  FAISS_ASSERT(outIndices.getSize(0) == queries.getSize(0));
211  FAISS_ASSERT(outDistances.getSize(1) == k);
212  FAISS_ASSERT(outIndices.getSize(1) == k);
213 
214  auto& mem = resources->getMemoryManagerCurrentDevice();
215  auto defaultStream = resources->getDefaultStreamCurrentDevice();
216 
217  // If we're quering against a 0 sized set, just return empty results
218  if (centroids.numElements() == 0) {
219  thrust::fill(thrust::cuda::par.on(defaultStream),
220  outDistances.data(), outDistances.end(),
221  Limits<T>::getMax());
222 
223  thrust::fill(thrust::cuda::par.on(defaultStream),
224  outIndices.data(), outIndices.end(),
225  -1);
226 
227  return;
228  }
229 
230  //
231  // Handle the problem in row tiles, to avoid excessive temporary
232  // memory requests
233  //
234 
235  FAISS_ASSERT(k <= centroids.getSize(0));
236  FAISS_ASSERT(k <= 1024); // select limitation
237 
238  int tileSize =
239  chooseTileSize<T>(
240  tileSizeOverride,
241  centroids.getSize(0),
242  resources->getMemoryManagerCurrentDevice().getSizeAvailable());
243 
244  int maxQueriesPerIteration = std::min(tileSize, queries.getSize(0));
245 
246  // Temporary output memory space we'll use
247  DeviceTensor<T, 2, true> distanceBuf1(
248  mem, {maxQueriesPerIteration, centroids.getSize(0)}, defaultStream);
249  DeviceTensor<T, 2, true> distanceBuf2(
250  mem, {maxQueriesPerIteration, centroids.getSize(0)}, defaultStream);
251  DeviceTensor<T, 2, true>* distanceBufs[2] =
252  {&distanceBuf1, &distanceBuf2};
253 
254  auto streams = resources->getAlternateStreamsCurrentDevice();
255  streamWait(streams, {defaultStream});
256 
257  int curStream = 0;
258 
259  for (int i = 0; i < queries.getSize(0); i += maxQueriesPerIteration) {
260  int numQueriesForIteration = std::min(maxQueriesPerIteration,
261  queries.getSize(0) - i);
262 
263  auto distanceBufView =
264  distanceBufs[curStream]->narrowOutermost(0, numQueriesForIteration);
265  auto queryView =
266  queries.narrowOutermost(i, numQueriesForIteration);
267  auto outDistanceView =
268  outDistances.narrowOutermost(i, numQueriesForIteration);
269  auto outIndexView =
270  outIndices.narrowOutermost(i, numQueriesForIteration);
271 
272  // (query id x dim) x (centroid id, dim)' = (query id, centroid id)
273  runMatrixMult(distanceBufView, false,
274  queryView, false,
275  centroidsTransposed ? *centroidsTransposed : centroids,
276  centroidsTransposed ? false : true,
277  1.0f, 0.0f,
278  resources->getBlasHandleCurrentDevice(),
279  streams[curStream]);
280 
281  // top-k of dot products
282  // (query id, top k centroids)
283  runBlockSelect(distanceBufView,
284  outDistanceView,
285  outIndexView,
286  true, k, streams[curStream]);
287 
288  curStream = (curStream + 1) % 2;
289  }
290 
291  streamWait({defaultStream}, streams);
292 }
293 
294 //
295 // Instantiations of the distance templates
296 //
297 
298 void
299 runIPDistance(GpuResources* resources,
300  Tensor<float, 2, true>& vectors,
301  Tensor<float, 2, true>* vectorsTransposed,
302  Tensor<float, 2, true>& queries,
303  int k,
304  Tensor<float, 2, true>& outDistances,
305  Tensor<int, 2, true>& outIndices,
306  int tileSizeOverride) {
307  runIPDistance<float>(resources,
308  vectors,
309  vectorsTransposed,
310  queries,
311  k,
312  outDistances,
313  outIndices,
314  tileSizeOverride);
315 }
316 
317 #ifdef FAISS_USE_FLOAT16
318 void
319 runIPDistance(GpuResources* resources,
320  Tensor<half, 2, true>& vectors,
321  Tensor<half, 2, true>* vectorsTransposed,
322  Tensor<half, 2, true>& queries,
323  int k,
324  Tensor<half, 2, true>& outDistances,
325  Tensor<int, 2, true>& outIndices,
326  int tileSizeOverride) {
327  runIPDistance<half>(resources,
328  vectors,
329  vectorsTransposed,
330  queries,
331  k,
332  outDistances,
333  outIndices,
334  tileSizeOverride);
335 }
336 #endif
337 
338 void
339 runL2Distance(GpuResources* resources,
340  Tensor<float, 2, true>& vectors,
341  Tensor<float, 2, true>* vectorsTransposed,
342  Tensor<float, 1, true>* vectorNorms,
343  Tensor<float, 2, true>& queries,
344  int k,
345  Tensor<float, 2, true>& outDistances,
346  Tensor<int, 2, true>& outIndices,
347  bool ignoreOutDistances,
348  int tileSizeOverride) {
349  runL2Distance<float>(resources,
350  vectors,
351  vectorsTransposed,
352  vectorNorms,
353  queries,
354  k,
355  outDistances,
356  outIndices,
357  ignoreOutDistances,
358  tileSizeOverride);
359 }
360 
361 #ifdef FAISS_USE_FLOAT16
362 void
363 runL2Distance(GpuResources* resources,
364  Tensor<half, 2, true>& vectors,
365  Tensor<half, 2, true>* vectorsTransposed,
366  Tensor<half, 1, true>* vectorNorms,
367  Tensor<half, 2, true>& queries,
368  int k,
369  Tensor<half, 2, true>& outDistances,
370  Tensor<int, 2, true>& outIndices,
371  bool ignoreOutDistances,
372  int tileSizeOverride) {
373  runL2Distance<half>(resources,
374  vectors,
375  vectorsTransposed,
376  vectorNorms,
377  queries,
378  k,
379  outDistances,
380  outIndices,
381  ignoreOutDistances,
382  tileSizeOverride);
383 }
384 #endif
385 
386 } } // namespace
Generated by   doxygen 1.8.5