mirror of
https://gitlab.com/libeigen/eigen.git
synced 2026-04-10 11:34:33 +08:00
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.
This commit is contained in:
committed by
Antonio Sánchez
parent
41d5d5334b
commit
17268b155d
@@ -69,7 +69,7 @@ struct bfloat16 : public bfloat16_impl::bfloat16_base {
|
||||
template<class T>
|
||||
explicit EIGEN_DEVICE_FUNC EIGEN_CONSTEXPR bfloat16(const T& val)
|
||||
: bfloat16_impl::bfloat16_base(bfloat16_impl::float_to_bfloat16_rtne<internal::is_integral<T>::value>(static_cast<float>(val))) {}
|
||||
|
||||
|
||||
explicit EIGEN_DEVICE_FUNC bfloat16(float f)
|
||||
: bfloat16_impl::bfloat16_base(bfloat16_impl::float_to_bfloat16_rtne<false>(f)) {}
|
||||
|
||||
@@ -88,7 +88,7 @@ struct bfloat16 : public bfloat16_impl::bfloat16_base {
|
||||
// +0.0 and -0.0 become false, everything else becomes true.
|
||||
return (value & 0x7fff) != 0;
|
||||
}
|
||||
#endif
|
||||
#endif
|
||||
|
||||
};
|
||||
} // namespace Eigen
|
||||
@@ -272,10 +272,14 @@ EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC __bfloat16_raw truncate_to_bfloat16(const
|
||||
return output;
|
||||
}
|
||||
|
||||
EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC EIGEN_CONSTEXPR __bfloat16_raw raw_uint16_to_bfloat16(unsigned short value) {
|
||||
EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC EIGEN_CONSTEXPR __bfloat16_raw raw_uint16_to_bfloat16(numext::uint16_t value) {
|
||||
return __bfloat16_raw(value);
|
||||
}
|
||||
|
||||
EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC EIGEN_CONSTEXPR numext::uint16_t raw_bfloat16_as_uint16(const __bfloat16_raw& bf) {
|
||||
return bf.value;
|
||||
}
|
||||
|
||||
// float_to_bfloat16_rtne template specialization that does not make any
|
||||
// assumption about the value of its function argument (ff).
|
||||
template <>
|
||||
@@ -454,7 +458,7 @@ EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC __bfloat16_raw float_to_bfloat16_rtne<fals
|
||||
// float_to_bfloat16_rtne template specialization that assumes that its function
|
||||
// argument (ff) is either a normal floating point number, or +/-infinity, or
|
||||
// zero. Used to improve the runtime performance of conversion from an integer
|
||||
// type to bfloat16.
|
||||
// type to bfloat16.
|
||||
template <>
|
||||
EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC __bfloat16_raw float_to_bfloat16_rtne<true>(float ff) {
|
||||
#if (defined(EIGEN_HAS_CUDA_BF16) && defined(EIGEN_HAS_HIP_BF16))
|
||||
@@ -691,7 +695,17 @@ bool (isfinite)(const Eigen::bfloat16& h) {
|
||||
return (bfloat16_impl::isfinite)(h);
|
||||
}
|
||||
|
||||
} // namespace numext
|
||||
template <>
|
||||
EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC Eigen::bfloat16 bit_cast<Eigen::bfloat16, uint16_t>(const uint16_t& src) {
|
||||
return Eigen::bfloat16(Eigen::bfloat16_impl::raw_uint16_to_bfloat16(src));
|
||||
}
|
||||
|
||||
template <>
|
||||
EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC uint16_t bit_cast<uint16_t, Eigen::bfloat16>(const Eigen::bfloat16& src) {
|
||||
return Eigen::bfloat16_impl::raw_bfloat16_as_uint16(src);
|
||||
}
|
||||
|
||||
} // namespace numext
|
||||
} // namespace Eigen
|
||||
|
||||
#endif // EIGEN_BFLOAT16_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);
|
||||
}
|
||||
} // namespace numext
|
||||
} // namespace Eigen
|
||||
|
||||
#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 // EIGEN_HALF_H
|
||||
|
||||
Reference in New Issue
Block a user