Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

All Classes Namespaces Functions Variables Typedefs Enumerations Enumerator Friends
Float16.cu
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 #include "Float16.cuh"
10 #include "nvidia/fp16_emu.cuh"
11 #include <thrust/execution_policy.h>
12 #include <thrust/transform.h>
13
14 #ifdef FAISS_USE_FLOAT16
15
16 namespace faiss { namespace gpu {
17
18 bool getDeviceSupportsFloat16Math(int device) {
19 const auto& prop = getDeviceProperties(device);
20
21 return (prop.major >= 6 ||
22 (prop.major == 5 && prop.minor >= 3));
23 }
24
25 struct FloatToHalf {
26 __device__ half operator()(float v) const { return __float2half(v); }
27 };
28
29 struct HalfToFloat {
30 __device__ float operator()(half v) const { return __half2float(v); }
31 };
32
33 void runConvertToFloat16(half* out,
34 const float* in,
35 size_t num,
36 cudaStream_t stream) {
37 thrust::transform(thrust::cuda::par.on(stream),
38 in, in + num, out, FloatToHalf());
39 }
40
41 void runConvertToFloat32(float* out,
42 const half* in,
43 size_t num,
44 cudaStream_t stream) {
45 thrust::transform(thrust::cuda::par.on(stream),
46 in, in + num, out, HalfToFloat());
47 }
48
49 __half hostFloat2Half(float a) {
50 #if CUDA_VERSION >= 9000
51 __half_raw raw;
52 raw.x = cpu_float2half_rn(a).x;
53 return __half(raw);
54 #else
55 __half h;
56 h.x = cpu_float2half_rn(a).x;
57 return h;
58 #endif
59 }
60
61 } } // namespace
62
63 #endif // FAISS_USE_FLOAT16