Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
IVFUtils.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 "IVFUtils.cuh"
13 #include "../utils/DeviceUtils.h"
14 #include "../utils/StaticUtils.h"
15 #include "../utils/Tensor.cuh"
16 #include "../utils/ThrustAllocator.cuh"
17 #include <thrust/scan.h>
18 #include <thrust/execution_policy.h>
19
20 namespace faiss { namespace gpu {
21
22 // Calculates the total number of intermediate distances to consider
23 // for all queries
24 __global__ void
25 getResultLengths(Tensor<int, 2, true> topQueryToCentroid,
26 int* listLengths,
27 int totalSize,
28 Tensor<int, 2, true> length) {
29 int linearThreadId = blockIdx.x * blockDim.x + threadIdx.x;
30 if (linearThreadId >= totalSize) {
31 return;
32 }
33
34 int nprobe = topQueryToCentroid.getSize(1);
35 int queryId = linearThreadId / nprobe;
36 int listId = linearThreadId % nprobe;
37
38 int centroidId = topQueryToCentroid[queryId][listId];
39
40 // Safety guard in case NaNs in input cause no list ID to be generated
41 length[queryId][listId] = (centroidId != -1) ? listLengths[centroidId] : 0;
42 }
43
44 void runCalcListOffsets(Tensor<int, 2, true>& topQueryToCentroid,
45 thrust::device_vector<int>& listLengths,
46 Tensor<int, 2, true>& prefixSumOffsets,
47 Tensor<char, 1, true>& thrustMem,
48 cudaStream_t stream) {
49 FAISS_ASSERT(topQueryToCentroid.getSize(0) == prefixSumOffsets.getSize(0));
50 FAISS_ASSERT(topQueryToCentroid.getSize(1) == prefixSumOffsets.getSize(1));
51
52 int totalSize = topQueryToCentroid.numElements();
53
54 int numThreads = std::min(totalSize, getMaxThreadsCurrentDevice());
55 int numBlocks = utils::divUp(totalSize, numThreads);
56
57 auto grid = dim3(numBlocks);
58 auto block = dim3(numThreads);
59
60 getResultLengths<<<grid, block, 0, stream>>>(
61 topQueryToCentroid,
62 listLengths.data().get(),
63 totalSize,
64 prefixSumOffsets);
65
66 // Prefix sum of the indices, so we know where the intermediate
67 // results should be maintained
68 // Thrust wants a place for its temporary allocations, so provide
69 // one, so it won't call cudaMalloc/Free
70 GpuResourcesThrustAllocator alloc(thrustMem.data(),
71 thrustMem.getSizeInBytes());
72
73 thrust::inclusive_scan(thrust::cuda::par(alloc).on(stream),
74 prefixSumOffsets.data(),
75 prefixSumOffsets.data() + totalSize,
76 prefixSumOffsets.data());
77 }
78
79 } } // namespace