add nextafter for bfloat16

This commit is contained in:
Peter Gavin
2024-10-21 21:23:41 +00:00
parent 53b83cddf9
commit b15ebb1c2d
2 changed files with 60 additions and 0 deletions

View File

@@ -763,6 +763,31 @@ EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC uint16_t bit_cast<uint16_t, Eigen::bfloat1
return Eigen::bfloat16_impl::raw_bfloat16_as_uint16(src);
}
EIGEN_STRONG_INLINE EIGEN_DEVICE_FUNC bfloat16 nextafter(const bfloat16& from, const bfloat16& to) {
if (numext::isnan EIGEN_NOT_A_MACRO(from)) {
return from;
}
if (numext::isnan EIGEN_NOT_A_MACRO(to)) {
return to;
}
if (from == to) {
return to;
}
uint16_t from_bits = numext::bit_cast<uint16_t>(from);
bool from_sign = from_bits >> 15;
// Whether we are adjusting toward the infinity with the same sign as from.
bool toward_inf = (to > from) == !from_sign;
if (toward_inf) {
++from_bits;
} else if ((from_bits & 0x7fff) == 0) {
// Adjusting away from inf, but from is zero, so just toggle the sign.
from_bits ^= 0x8000;
} else {
--from_bits;
}
return numext::bit_cast<bfloat16>(from_bits);
}
} // namespace numext
} // namespace Eigen