diff options
author | Antonio Sanchez <cantonios@google.com> | 2020-11-17 15:32:44 -0800 |
---|---|---|
committer | Antonio Sánchez <cantonios@google.com> | 2020-11-18 20:32:35 +0000 |
commit | 17268b155d54422f1294130c0fb8c178757d911a (patch) | |
tree | 2be3d541729f3e9be6a180a58270bae10156df4f /Eigen/src/Core/arch/Default/Half.h | |
parent | 41d5d5334b8a4e364dfd88dcd91f6cd38834b8ed (diff) |
Add bit_cast for half/bfloat to/from uint16_t, fix TensorRandom
The existing `TensorRandom.h` implementation makes the assumption that
`half` (`bfloat16`) has a `uint16_t` member `x` (`value`), which is not
always true. This currently fails on arm64, where `x` has type `__fp16`.
Added `bit_cast` specializations to allow casting to/from `uint16_t`
for both `half` and `bfloat16`. Also added tests in
`half_float`, `bfloat16_float`, and `cxx11_tensor_random` to catch
these errors in the future.
Diffstat (limited to 'Eigen/src/Core/arch/Default/Half.h')
-rw-r--r-- | Eigen/src/Core/arch/Default/Half.h | 30 |
1 files changed, 28 insertions, 2 deletions
diff --git a/Eigen/src/Core/arch/Default/Half.h b/Eigen/src/Core/arch/Default/Half.h index 0bc1e9d19..fda38bcb0 100644 --- a/Eigen/src/Core/arch/Default/Half.h +++ b/Eigen/src/Core/arch/Default/Half.h @@ -494,6 +494,19 @@ EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC EIGEN_CONSTEXPR __half_raw raw_uint16_to_h #endif } +EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC EIGEN_CONSTEXPR numext::uint16_t raw_half_as_uint16(const __half_raw& h) { + // HIP/CUDA/Default have a member 'x' of type uint16_t. + // For ARM64 native half, the member 'x' is of type __fp16, so we need to bit-cast. + // For SYCL, cl::sycl::half is _Float16, so cast directly. +#if defined(EIGEN_HAS_ARM64_FP16_SCALAR_ARITHMETIC) + return numext::bit_cast<numext::uint16_t>(h.x); +#elif defined(SYCL_DEVICE_ONLY) + return numext::bit_cast<numext::uint16_t>(h); +#else + return h.x; +#endif +} + union float32_bits { unsigned int u; float f; @@ -812,10 +825,11 @@ EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Eigen::half __ldg(const Eigen::half* ptr) } #endif -#if defined(EIGEN_GPU_COMPILE_PHASE) namespace Eigen { namespace numext { +#if defined(EIGEN_GPU_COMPILE_PHASE) + template <> EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE bool(isnan)(const Eigen::half& h) { return (half_impl::isnan)(h); @@ -830,8 +844,20 @@ template <> EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE bool(isfinite)(const Eigen::half& h) { return (half_impl::isfinite)(h); } + +#endif + +template <> +EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Eigen::half bit_cast<Eigen::half, uint16_t>(const uint16_t& src) { + return Eigen::half(Eigen::half_impl::raw_uint16_to_half(src)); +} + +template <> +EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC uint16_t bit_cast<uint16_t, Eigen::half>(const Eigen::half& src) { + return Eigen::half_impl::raw_half_as_uint16(src); +} + } // namespace numext } // namespace Eigen -#endif #endif // EIGEN_HALF_H |