Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
VectorResidual.cu
1 /**
2 * Copyright (c) 2015-present, Facebook, Inc.
3 * All rights reserved.
4 *
5 * This source code is licensed under the BSD+Patents license found in the
6 * LICENSE file in the root directory of this source tree.
7 */
8
9 #include "VectorResidual.cuh"
10 #include "../../FaissAssert.h"
11 #include "../utils/ConversionOperators.cuh"
12 #include "../utils/DeviceUtils.h"
13 #include "../utils/Tensor.cuh"
14 #include "../utils/StaticUtils.h"
15 #include <math_constants.h> // in CUDA SDK, for CUDART_NAN_F
16
17 namespace faiss { namespace gpu {
18
19 template <typename CentroidT, bool LargeDim>
20 __global__ void calcResidual(Tensor<float, 2, true> vecs,
21 Tensor<CentroidT, 2, true> centroids,
22 Tensor<int, 1, true> vecToCentroid,
23 Tensor<float, 2, true> residuals) {
24 auto vec = vecs[blockIdx.x];
25 auto residual = residuals[blockIdx.x];
26
27 int centroidId = vecToCentroid[blockIdx.x];
28 // Vector could be invalid (containing NaNs), so -1 was the
29 // classified centroid
30 if (centroidId == -1) {
31 if (LargeDim) {
32 for (int i = threadIdx.x; i < vecs.getSize(1); i += blockDim.x) {
33 residual[i] = CUDART_NAN_F;
34 }
35 } else {
36 residual[threadIdx.x] = CUDART_NAN_F;
37 }
38
39 return;
40 }
41
42 auto centroid = centroids[centroidId];
43
44 if (LargeDim) {
45 for (int i = threadIdx.x; i < vecs.getSize(1); i += blockDim.x) {
46 residual[i] = vec[i] - ConvertTo<float>::to(centroid[i]);
47 }
48 } else {
49 residual[threadIdx.x] = vec[threadIdx.x] -
50 ConvertTo<float>::to(centroid[threadIdx.x]);
51 }
52 }
53
54 template <typename CentroidT>
55 void calcResidual(Tensor<float, 2, true>& vecs,
56 Tensor<CentroidT, 2, true>& centroids,
57 Tensor<int, 1, true>& vecToCentroid,
58 Tensor<float, 2, true>& residuals,
59 cudaStream_t stream) {
60 FAISS_ASSERT(vecs.getSize(1) == centroids.getSize(1));
61 FAISS_ASSERT(vecs.getSize(1) == residuals.getSize(1));
62 FAISS_ASSERT(vecs.getSize(0) == vecToCentroid.getSize(0));
63 FAISS_ASSERT(vecs.getSize(0) == residuals.getSize(0));
64
65 dim3 grid(vecs.getSize(0));
66
67 int maxThreads = getMaxThreadsCurrentDevice();
68 bool largeDim = vecs.getSize(1) > maxThreads;
69 dim3 block(std::min(vecs.getSize(1), maxThreads));
70
71 if (largeDim) {
72 calcResidual<CentroidT, true><<<grid, block, 0, stream>>>(
73 vecs, centroids, vecToCentroid, residuals);
74 } else {
75 calcResidual<CentroidT, false><<<grid, block, 0, stream>>>(
76 vecs, centroids, vecToCentroid, residuals);
77 }
78
79 CUDA_TEST_ERROR();
80 }
81
82 void runCalcResidual(Tensor<float, 2, true>& vecs,
83 Tensor<float, 2, true>& centroids,
84 Tensor<int, 1, true>& vecToCentroid,
85 Tensor<float, 2, true>& residuals,
86 cudaStream_t stream) {
87 calcResidual<float>(vecs, centroids, vecToCentroid, residuals, stream);
88 }
89
90 #ifdef FAISS_USE_FLOAT16
91 void runCalcResidual(Tensor<float, 2, true>& vecs,
92 Tensor<half, 2, true>& centroids,
93 Tensor<int, 1, true>& vecToCentroid,
94 Tensor<float, 2, true>& residuals,
95 cudaStream_t stream) {
96 calcResidual<half>(vecs, centroids, vecToCentroid, residuals, stream);
97 }
98 #endif
99
100 } } // namespace