1// This file is part of Eigen, a lightweight C++ template library
2// for linear algebra.
3//
4// Copyright (C) 2016 Benoit Steiner <benoit.steiner.goog@gmail.com>
5//
6// This Source Code Form is subject to the terms of the Mozilla
7// Public License v. 2.0. If a copy of the MPL was not distributed
8// with this file, You can obtain one at http://mozilla.org/MPL/2.0/.
9
10#ifndef EIGEN_TYPE_CASTING_CUDA_H
11#define EIGEN_TYPE_CASTING_CUDA_H
12
13namespace Eigen {
14
15namespace internal {
16
17template<>
18struct scalar_cast_op<float, Eigen::half> {
19 EIGEN_EMPTY_STRUCT_CTOR(scalar_cast_op)
20 typedef Eigen::half result_type;
21 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Eigen::half operator() (const float& a) const {
22 #if defined(EIGEN_HAS_CUDA_FP16) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 300
23 return __float2half(a);
24 #else
25 return Eigen::half(a);
26 #endif
27 }
28};
29
30template<>
31struct functor_traits<scalar_cast_op<float, Eigen::half> >
32{ enum { Cost = NumTraits<float>::AddCost, PacketAccess = false }; };
33
34
35template<>
36struct scalar_cast_op<int, Eigen::half> {
37 EIGEN_EMPTY_STRUCT_CTOR(scalar_cast_op)
38 typedef Eigen::half result_type;
39 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Eigen::half operator() (const int& a) const {
40 #if defined(EIGEN_HAS_CUDA_FP16) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 300
41 return __float2half(static_cast<float>(a));
42 #else
43 return Eigen::half(static_cast<float>(a));
44 #endif
45 }
46};
47
48template<>
49struct functor_traits<scalar_cast_op<int, Eigen::half> >
50{ enum { Cost = NumTraits<float>::AddCost, PacketAccess = false }; };
51
52
53template<>
54struct scalar_cast_op<Eigen::half, float> {
55 EIGEN_EMPTY_STRUCT_CTOR(scalar_cast_op)
56 typedef float result_type;
57 EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE float operator() (const Eigen::half& a) const {
58 #if defined(EIGEN_HAS_CUDA_FP16) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 300
59 return __half2float(a);
60 #else
61 return static_cast<float>(a);
62 #endif
63 }
64};
65
66template<>
67struct functor_traits<scalar_cast_op<Eigen::half, float> >
68{ enum { Cost = NumTraits<float>::AddCost, PacketAccess = false }; };
69
70
71
72#if defined(EIGEN_HAS_CUDA_FP16) && defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 300
73
74template <>
75struct type_casting_traits<Eigen::half, float> {
76 enum {
77 VectorizedCast = 1,
78 SrcCoeffRatio = 2,
79 TgtCoeffRatio = 1
80 };
81};
82
83template<> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE float4 pcast<half2, float4>(const half2& a, const half2& b) {
84 float2 r1 = __half22float2(a);
85 float2 r2 = __half22float2(b);
86 return make_float4(r1.x, r1.y, r2.x, r2.y);
87}
88
89template <>
90struct type_casting_traits<float, Eigen::half> {
91 enum {
92 VectorizedCast = 1,
93 SrcCoeffRatio = 1,
94 TgtCoeffRatio = 2
95 };
96};
97
98template<> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE half2 pcast<float4, half2>(const float4& a) {
99 // Simply discard the second half of the input
100 return __floats2half2_rn(a.x, a.y);
101}
102
103#elif defined EIGEN_VECTORIZE_AVX512
104template <>
105struct type_casting_traits<half, float> {
106 enum {
107 VectorizedCast = 1,
108 SrcCoeffRatio = 1,
109 TgtCoeffRatio = 1
110 };
111};
112
113template<> EIGEN_STRONG_INLINE Packet16f pcast<Packet16h, Packet16f>(const Packet16h& a) {
114 return half2float(a);
115}
116
117template <>
118struct type_casting_traits<float, half> {
119 enum {
120 VectorizedCast = 1,
121 SrcCoeffRatio = 1,
122 TgtCoeffRatio = 1
123 };
124};
125
126template<> EIGEN_STRONG_INLINE Packet16h pcast<Packet16f, Packet16h>(const Packet16f& a) {
127 return float2half(a);
128}
129
130#elif defined EIGEN_VECTORIZE_AVX
131
132template <>
133struct type_casting_traits<Eigen::half, float> {
134 enum {
135 VectorizedCast = 1,
136 SrcCoeffRatio = 1,
137 TgtCoeffRatio = 1
138 };
139};
140
141template<> EIGEN_STRONG_INLINE Packet8f pcast<Packet8h, Packet8f>(const Packet8h& a) {
142 return half2float(a);
143}
144
145template <>
146struct type_casting_traits<float, Eigen::half> {
147 enum {
148 VectorizedCast = 1,
149 SrcCoeffRatio = 1,
150 TgtCoeffRatio = 1
151 };
152};
153
154template<> EIGEN_STRONG_INLINE Packet8h pcast<Packet8f, Packet8h>(const Packet8f& a) {
155 return float2half(a);
156}
157
158// Disable the following code since it's broken on too many platforms / compilers.
159//#elif defined(EIGEN_VECTORIZE_SSE) && (!EIGEN_ARCH_x86_64) && (!EIGEN_COMP_MSVC)
160#elif 0
161
162template <>
163struct type_casting_traits<Eigen::half, float> {
164 enum {
165 VectorizedCast = 1,
166 SrcCoeffRatio = 1,
167 TgtCoeffRatio = 1
168 };
169};
170
171template<> EIGEN_STRONG_INLINE Packet4f pcast<Packet4h, Packet4f>(const Packet4h& a) {
172 __int64_t a64 = _mm_cvtm64_si64(a.x);
173 Eigen::half h = raw_uint16_to_half(static_cast<unsigned short>(a64));
174 float f1 = static_cast<float>(h);
175 h = raw_uint16_to_half(static_cast<unsigned short>(a64 >> 16));
176 float f2 = static_cast<float>(h);
177 h = raw_uint16_to_half(static_cast<unsigned short>(a64 >> 32));
178 float f3 = static_cast<float>(h);
179 h = raw_uint16_to_half(static_cast<unsigned short>(a64 >> 48));
180 float f4 = static_cast<float>(h);
181 return _mm_set_ps(f4, f3, f2, f1);
182}
183
184template <>
185struct type_casting_traits<float, Eigen::half> {
186 enum {
187 VectorizedCast = 1,
188 SrcCoeffRatio = 1,
189 TgtCoeffRatio = 1
190 };
191};
192
193template<> EIGEN_STRONG_INLINE Packet4h pcast<Packet4f, Packet4h>(const Packet4f& a) {
194 EIGEN_ALIGN16 float aux[4];
195 pstore(aux, a);
196 Eigen::half h0(aux[0]);
197 Eigen::half h1(aux[1]);
198 Eigen::half h2(aux[2]);
199 Eigen::half h3(aux[3]);
200
201 Packet4h result;
202 result.x = _mm_set_pi16(h3.x, h2.x, h1.x, h0.x);
203 return result;
204}
205
206#endif
207
208} // end namespace internal
209
210} // end namespace Eigen
211
212#endif // EIGEN_TYPE_CASTING_CUDA_H
213