Faiss
 All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
Distance.cu
1 /**
2  * Copyright (c) 2015-present, Facebook, Inc.
3  * All rights reserved.
4  *
5  * This source code is licensed under the CC-by-NC license found in the
6  * LICENSE file in the root directory of this source tree.
7  */
8 
9 // Copyright 2004-present Facebook. All Rights Reserved.
10 
11 #include "Distance.cuh"
12 #include "BroadcastSum.cuh"
13 #include "L2Norm.cuh"
14 #include "L2Select.cuh"
15 #include "../../FaissAssert.h"
16 #include "../GpuResources.h"
17 #include "../utils/DeviceUtils.h"
18 #include "../utils/Limits.cuh"
19 #include "../utils/MatrixMult.cuh"
20 #include "../utils/BlockSelectKernel.cuh"
21 
22 #include <memory>
23 #include <thrust/fill.h>
24 #include <thrust/for_each.h>
25 #include <thrust/device_ptr.h>
26 #include <thrust/execution_policy.h>
27 
28 namespace faiss { namespace gpu {
29 
30 namespace {
31 
32 constexpr int kDefaultTileSize = 256;
33 
34 template <typename T>
35 int chooseTileSize(int tileSizeOverride,
36  size_t numCentroids,
37  size_t tempMemAvailable) {
38  if (tileSizeOverride > 0) {
39  return tileSizeOverride;
40  }
41 
42  size_t tileSize =
43  sizeof(T) < 4 ? kDefaultTileSize * 2 : kDefaultTileSize;
44 
45  while (tileSize > 64) {
46  size_t memRequirement = 2 * tileSize * numCentroids * sizeof(T);
47 
48  if (memRequirement <= tempMemAvailable) {
49  // This fits entirely into our temporary memory
50  return tileSize;
51  }
52 
53  // Otherwise, halve the tile size
54  tileSize /= 2;
55  }
56 
57  // We use 64 as the minimum acceptable tile size
58  FAISS_ASSERT(tileSize >= 64);
59 
60  // FIXME: if we're running with no available temp memory, do we try
61  // and go larger based on free memory available on the device?
62 
63  return tileSize;
64 }
65 
66 }
67 
68 template <typename T>
69 void runL2Distance(GpuResources* resources,
70  Tensor<T, 2, true>& centroids,
71  Tensor<T, 2, true>* centroidsTransposed,
72  Tensor<T, 1, true>* centroidNorms,
73  Tensor<T, 2, true>& queries,
74  int k,
75  Tensor<T, 2, true>& outDistances,
76  Tensor<int, 2, true>& outIndices,
77  bool useHgemm,
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, useHgemm,
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  bool useHgemm,
209  int tileSizeOverride = -1) {
210  FAISS_ASSERT(outDistances.getSize(0) == queries.getSize(0));
211  FAISS_ASSERT(outIndices.getSize(0) == queries.getSize(0));
212  FAISS_ASSERT(outDistances.getSize(1) == k);
213  FAISS_ASSERT(outIndices.getSize(1) == k);
214 
215  auto& mem = resources->getMemoryManagerCurrentDevice();
216  auto defaultStream = resources->getDefaultStreamCurrentDevice();
217 
218  // If we're quering against a 0 sized set, just return empty results
219  if (centroids.numElements() == 0) {
220  thrust::fill(thrust::cuda::par.on(defaultStream),
221  outDistances.data(), outDistances.end(),
222  Limits<T>::getMax());
223 
224  thrust::fill(thrust::cuda::par.on(defaultStream),
225  outIndices.data(), outIndices.end(),
226  -1);
227 
228  return;
229  }
230 
231  //
232  // Handle the problem in row tiles, to avoid excessive temporary
233  // memory requests
234  //
235 
236  FAISS_ASSERT(k <= centroids.getSize(0));
237  FAISS_ASSERT(k <= 1024); // select limitation
238 
239  int tileSize =
240  chooseTileSize<T>(
241  tileSizeOverride,
242  centroids.getSize(0),
243  resources->getMemoryManagerCurrentDevice().getSizeAvailable());
244 
245  int maxQueriesPerIteration = std::min(tileSize, queries.getSize(0));
246 
247  // Temporary output memory space we'll use
248  DeviceTensor<T, 2, true> distanceBuf1(
249  mem, {maxQueriesPerIteration, centroids.getSize(0)}, defaultStream);
250  DeviceTensor<T, 2, true> distanceBuf2(
251  mem, {maxQueriesPerIteration, centroids.getSize(0)}, defaultStream);
252  DeviceTensor<T, 2, true>* distanceBufs[2] =
253  {&distanceBuf1, &distanceBuf2};
254 
255  auto streams = resources->getAlternateStreamsCurrentDevice();
256  streamWait(streams, {defaultStream});
257 
258  int curStream = 0;
259 
260  for (int i = 0; i < queries.getSize(0); i += maxQueriesPerIteration) {
261  int numQueriesForIteration = std::min(maxQueriesPerIteration,
262  queries.getSize(0) - i);
263 
264  auto distanceBufView =
265  distanceBufs[curStream]->narrowOutermost(0, numQueriesForIteration);
266  auto queryView =
267  queries.narrowOutermost(i, numQueriesForIteration);
268  auto outDistanceView =
269  outDistances.narrowOutermost(i, numQueriesForIteration);
270  auto outIndexView =
271  outIndices.narrowOutermost(i, numQueriesForIteration);
272 
273  // (query id x dim) x (centroid id, dim)' = (query id, centroid id)
274  runMatrixMult(distanceBufView, false,
275  queryView, false,
276  centroidsTransposed ? *centroidsTransposed : centroids,
277  centroidsTransposed ? false : true,
278  1.0f, 0.0f, useHgemm,
279  resources->getBlasHandleCurrentDevice(),
280  streams[curStream]);
281 
282  // top-k of dot products
283  // (query id, top k centroids)
284  runBlockSelect(distanceBufView,
285  outDistanceView,
286  outIndexView,
287  true, k, streams[curStream]);
288 
289  curStream = (curStream + 1) % 2;
290  }
291 
292  streamWait({defaultStream}, streams);
293 }
294 
295 //
296 // Instantiations of the distance templates
297 //
298 
299 void
300 runIPDistance(GpuResources* resources,
301  Tensor<float, 2, true>& vectors,
302  Tensor<float, 2, true>* vectorsTransposed,
303  Tensor<float, 2, true>& queries,
304  int k,
305  Tensor<float, 2, true>& outDistances,
306  Tensor<int, 2, true>& outIndices,
307  int tileSizeOverride) {
308  runIPDistance<float>(resources,
309  vectors,
310  vectorsTransposed,
311  queries,
312  k,
313  outDistances,
314  outIndices,
315  false,
316  tileSizeOverride);
317 }
318 
319 #ifdef FAISS_USE_FLOAT16
320 void
321 runIPDistance(GpuResources* resources,
322  Tensor<half, 2, true>& vectors,
323  Tensor<half, 2, true>* vectorsTransposed,
324  Tensor<half, 2, true>& queries,
325  int k,
326  Tensor<half, 2, true>& outDistances,
327  Tensor<int, 2, true>& outIndices,
328  bool useHgemm,
329  int tileSizeOverride) {
330  runIPDistance<half>(resources,
331  vectors,
332  vectorsTransposed,
333  queries,
334  k,
335  outDistances,
336  outIndices,
337  useHgemm,
338  tileSizeOverride);
339 }
340 #endif
341 
342 void
343 runL2Distance(GpuResources* resources,
344  Tensor<float, 2, true>& vectors,
345  Tensor<float, 2, true>* vectorsTransposed,
346  Tensor<float, 1, true>* vectorNorms,
347  Tensor<float, 2, true>& queries,
348  int k,
349  Tensor<float, 2, true>& outDistances,
350  Tensor<int, 2, true>& outIndices,
351  bool ignoreOutDistances,
352  int tileSizeOverride) {
353  runL2Distance<float>(resources,
354  vectors,
355  vectorsTransposed,
356  vectorNorms,
357  queries,
358  k,
359  outDistances,
360  outIndices,
361  false,
362  ignoreOutDistances,
363  tileSizeOverride);
364 }
365 
366 #ifdef FAISS_USE_FLOAT16
367 void
368 runL2Distance(GpuResources* resources,
369  Tensor<half, 2, true>& vectors,
370  Tensor<half, 2, true>* vectorsTransposed,
371  Tensor<half, 1, true>* vectorNorms,
372  Tensor<half, 2, true>& queries,
373  int k,
374  Tensor<half, 2, true>& outDistances,
375  Tensor<int, 2, true>& outIndices,
376  bool useHgemm,
377  bool ignoreOutDistances,
378  int tileSizeOverride) {
379  runL2Distance<half>(resources,
380  vectors,
381  vectorsTransposed,
382  vectorNorms,
383  queries,
384  k,
385  outDistances,
386  outIndices,
387  useHgemm,
388  ignoreOutDistances,
389  tileSizeOverride);
390 }
391 #endif
392 
393 } } // namespace
Generated by   doxygen 1.8.5