Faiss

- Main Page
- Namespaces
- Classes
- Files

- File List

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