Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
PQCodeDistances.cuh
1 /**
2 * Copyright (c) Facebook, Inc. and its affiliates.
3 *
4 * This source code is licensed under the MIT license found in the
5 * LICENSE file in the root directory of this source tree.
6 */
7
8
9 #pragma once
10
11 #include "../utils/Tensor.cuh"
12 #include "../utils/NoTypeTensor.cuh"
13 #include <cublas_v2.h>
14
15 namespace faiss { namespace gpu {
16
17 class DeviceMemory;
18
19 /// pqCentroids is of the form (sub q)(sub dim)(code id)
20 /// Calculates the distance from the (query - centroid) residual to
21 /// each sub-code vector, for the given list of query results in
22 /// topQueryToCentroid
23 void runPQCodeDistances(Tensor<float, 3, true>& pqCentroids,
24 Tensor<float, 2, true>& queries,
25 Tensor<float, 2, true>& coarseCentroids,
26 Tensor<int, 2, true>& topQueryToCentroid,
27 NoTypeTensor<4, true>& outCodeDistances,
28 bool useFloat16Lookup,
29 cudaStream_t stream);
30
31 void runPQCodeDistancesMM(Tensor<float, 3, true>& pqCentroids,
32 Tensor<float, 2, true>& queries,
33 Tensor<float, 2, true>& coarseCentroids,
34 Tensor<int, 2, true>& topQueryToCentroid,
35 NoTypeTensor<4, true>& outCodeDistances,
36 bool useFloat16Lookup,
37 DeviceMemory& mem,
38 cublasHandle_t handle,
39 cudaStream_t stream);
40
41 } } // namespace