Special function implementations for half/bfloat16 packets.

Current implementations fail to consider half-float packets, only
half-float scalars.  Added specializations for packets on AVX, AVX512 and
NEON.  Added tests to `special_packetmath`.

The current `special_functions` tests would fail for half and bfloat16 due to
lack of precision. The NEON tests also fail with precision issues and
due to different handling of `sqrt(inf)`, so special functions bessel, ndtri
have been disabled.

Tested with AVX, AVX512.
This commit is contained in:
Antonio Sanchez
2020-12-02 14:00:57 -08:00
parent 305b8bd277
commit e2f21465fe
17 changed files with 418 additions and 159 deletions

View File

@@ -0,0 +1,46 @@
#ifndef EIGEN_AVX_BESSELFUNCTIONS_H
#define EIGEN_AVX_BESSELFUNCTIONS_H
namespace Eigen {
namespace internal {
F16_PACKET_FUNCTION(Packet8f, Packet8h, pbessel_i0)
BF16_PACKET_FUNCTION(Packet8f, Packet8bf, pbessel_i0)
F16_PACKET_FUNCTION(Packet8f, Packet8h, pbessel_i0e)
BF16_PACKET_FUNCTION(Packet8f, Packet8bf, pbessel_i0e)
F16_PACKET_FUNCTION(Packet8f, Packet8h, pbessel_i1)
BF16_PACKET_FUNCTION(Packet8f, Packet8bf, pbessel_i1)
F16_PACKET_FUNCTION(Packet8f, Packet8h, pbessel_i1e)
BF16_PACKET_FUNCTION(Packet8f, Packet8bf, pbessel_i1e)
F16_PACKET_FUNCTION(Packet8f, Packet8h, pbessel_j0)
BF16_PACKET_FUNCTION(Packet8f, Packet8bf, pbessel_j0)
F16_PACKET_FUNCTION(Packet8f, Packet8h, pbessel_j1)
BF16_PACKET_FUNCTION(Packet8f, Packet8bf, pbessel_j1)
F16_PACKET_FUNCTION(Packet8f, Packet8h, pbessel_k0)
BF16_PACKET_FUNCTION(Packet8f, Packet8bf, pbessel_k0)
F16_PACKET_FUNCTION(Packet8f, Packet8h, pbessel_k0e)
BF16_PACKET_FUNCTION(Packet8f, Packet8bf, pbessel_k0e)
F16_PACKET_FUNCTION(Packet8f, Packet8h, pbessel_k1)
BF16_PACKET_FUNCTION(Packet8f, Packet8bf, pbessel_k1)
F16_PACKET_FUNCTION(Packet8f, Packet8h, pbessel_k1e)
BF16_PACKET_FUNCTION(Packet8f, Packet8bf, pbessel_k1e)
F16_PACKET_FUNCTION(Packet8f, Packet8h, pbessel_y0)
BF16_PACKET_FUNCTION(Packet8f, Packet8bf, pbessel_y0)
F16_PACKET_FUNCTION(Packet8f, Packet8h, pbessel_y1)
BF16_PACKET_FUNCTION(Packet8f, Packet8bf, pbessel_y1)
} // namespace internal
} // namespace Eigen
#endif // EIGEN_AVX_BESSELFUNCTIONS_H

View File

@@ -0,0 +1,16 @@
#ifndef EIGEN_AVX_SPECIALFUNCTIONS_H
#define EIGEN_AVX_SPECIALFUNCTIONS_H
namespace Eigen {
namespace internal {
F16_PACKET_FUNCTION(Packet8f, Packet8h, perf)
BF16_PACKET_FUNCTION(Packet8f, Packet8bf, perf)
F16_PACKET_FUNCTION(Packet8f, Packet8h, pndtri)
BF16_PACKET_FUNCTION(Packet8f, Packet8bf, pndtri)
} // namespace internal
} // namespace Eigen
#endif // EIGEN_AVX_SPECIAL_FUNCTIONS_H

View File

@@ -0,0 +1,46 @@
#ifndef EIGEN_AVX512_BESSELFUNCTIONS_H
#define EIGEN_AVX512_BESSELFUNCTIONS_H
namespace Eigen {
namespace internal {
F16_PACKET_FUNCTION(Packet16f, Packet16h, pbessel_i0)
BF16_PACKET_FUNCTION(Packet16f, Packet16bf, pbessel_i0)
F16_PACKET_FUNCTION(Packet16f, Packet16h, pbessel_i0e)
BF16_PACKET_FUNCTION(Packet16f, Packet16bf, pbessel_i0e)
F16_PACKET_FUNCTION(Packet16f, Packet16h, pbessel_i1)
BF16_PACKET_FUNCTION(Packet16f, Packet16bf, pbessel_i1)
F16_PACKET_FUNCTION(Packet16f, Packet16h, pbessel_i1e)
BF16_PACKET_FUNCTION(Packet16f, Packet16bf, pbessel_i1e)
F16_PACKET_FUNCTION(Packet16f, Packet16h, pbessel_j0)
BF16_PACKET_FUNCTION(Packet16f, Packet16bf, pbessel_j0)
F16_PACKET_FUNCTION(Packet16f, Packet16h, pbessel_j1)
BF16_PACKET_FUNCTION(Packet16f, Packet16bf, pbessel_j1)
F16_PACKET_FUNCTION(Packet16f, Packet16h, pbessel_k0)
BF16_PACKET_FUNCTION(Packet16f, Packet16bf, pbessel_k0)
F16_PACKET_FUNCTION(Packet16f, Packet16h, pbessel_k0e)
BF16_PACKET_FUNCTION(Packet16f, Packet16bf, pbessel_k0e)
F16_PACKET_FUNCTION(Packet16f, Packet16h, pbessel_k1)
BF16_PACKET_FUNCTION(Packet16f, Packet16bf, pbessel_k1)
F16_PACKET_FUNCTION(Packet16f, Packet16h, pbessel_k1e)
BF16_PACKET_FUNCTION(Packet16f, Packet16bf, pbessel_k1e)
F16_PACKET_FUNCTION(Packet16f, Packet16h, pbessel_y0)
BF16_PACKET_FUNCTION(Packet16f, Packet16bf, pbessel_y0)
F16_PACKET_FUNCTION(Packet16f, Packet16h, pbessel_y1)
BF16_PACKET_FUNCTION(Packet16f, Packet16bf, pbessel_y1)
} // namespace internal
} // namespace Eigen
#endif // EIGEN_AVX512_BESSELFUNCTIONS_H

View File

@@ -0,0 +1,16 @@
#ifndef EIGEN_AVX512_SPECIALFUNCTIONS_H
#define EIGEN_AVX512_SPECIALFUNCTIONS_H
namespace Eigen {
namespace internal {
F16_PACKET_FUNCTION(Packet16f, Packet16h, perf)
BF16_PACKET_FUNCTION(Packet16f, Packet16bf, perf)
F16_PACKET_FUNCTION(Packet16f, Packet16h, pndtri)
BF16_PACKET_FUNCTION(Packet16f, Packet16bf, pndtri)
} // namespace internal
} // namespace Eigen
#endif // EIGEN_AVX512_SPECIAL_FUNCTIONS_H

View File

@@ -0,0 +1,54 @@
#ifndef EIGEN_NEON_BESSELFUNCTIONS_H
#define EIGEN_NEON_BESSELFUNCTIONS_H
namespace Eigen {
namespace internal {
#if EIGEN_HAS_ARM64_FP16_VECTOR_ARITHMETIC
#define NEON_HALF_TO_FLOAT_FUNCTIONS(METHOD) \
template <> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE \
Packet8hf METHOD<Packet8hf>(const Packet8hf& x) { \
const Packet4f lo = METHOD<Packet4f>(vcvt_f32_f16(vget_low_f16(x))); \
const Packet4f hi = METHOD<Packet4f>(vcvt_f32_f16(vget_high_f16(x))); \
return vcombine_f16(vcvt_f16_f32(lo), vcvt_f16_f32(hi)); \
} \
\
template <> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE \
Packet4hf METHOD<Packet4hf>(const Packet4hf& x) { \
return vcvt_f16_f32(METHOD<Packet4f>(vcvt_f32_f16(x))); \
}
NEON_HALF_TO_FLOAT_FUNCTIONS(pbessel_i0)
NEON_HALF_TO_FLOAT_FUNCTIONS(pbessel_i0e)
NEON_HALF_TO_FLOAT_FUNCTIONS(pbessel_i1)
NEON_HALF_TO_FLOAT_FUNCTIONS(pbessel_i1e)
NEON_HALF_TO_FLOAT_FUNCTIONS(pbessel_j0)
NEON_HALF_TO_FLOAT_FUNCTIONS(pbessel_j1)
NEON_HALF_TO_FLOAT_FUNCTIONS(pbessel_k0)
NEON_HALF_TO_FLOAT_FUNCTIONS(pbessel_k0e)
NEON_HALF_TO_FLOAT_FUNCTIONS(pbessel_k1)
NEON_HALF_TO_FLOAT_FUNCTIONS(pbessel_k1e)
NEON_HALF_TO_FLOAT_FUNCTIONS(pbessel_y0)
NEON_HALF_TO_FLOAT_FUNCTIONS(pbessel_y1)
#undef NEON_HALF_TO_FLOAT_FUNCTIONS
#endif
BF16_PACKET_FUNCTION(Packet4f, Packet4bf, pbessel_i0)
BF16_PACKET_FUNCTION(Packet4f, Packet4bf, pbessel_i0e)
BF16_PACKET_FUNCTION(Packet4f, Packet4bf, pbessel_i1)
BF16_PACKET_FUNCTION(Packet4f, Packet4bf, pbessel_i1e)
BF16_PACKET_FUNCTION(Packet4f, Packet4bf, pbessel_j0)
BF16_PACKET_FUNCTION(Packet4f, Packet4bf, pbessel_j1)
BF16_PACKET_FUNCTION(Packet4f, Packet4bf, pbessel_k0)
BF16_PACKET_FUNCTION(Packet4f, Packet4bf, pbessel_k0e)
BF16_PACKET_FUNCTION(Packet4f, Packet4bf, pbessel_k1)
BF16_PACKET_FUNCTION(Packet4f, Packet4bf, pbessel_k1e)
BF16_PACKET_FUNCTION(Packet4f, Packet4bf, pbessel_y0)
BF16_PACKET_FUNCTION(Packet4f, Packet4bf, pbessel_y1)
} // namespace internal
} // namespace Eigen
#endif // EIGEN_NEON_BESSELFUNCTIONS_H

View File

@@ -0,0 +1,34 @@
#ifndef EIGEN_NEON_SPECIALFUNCTIONS_H
#define EIGEN_NEON_SPECIALFUNCTIONS_H
namespace Eigen {
namespace internal {
#ifdef EIGEN_HAS_ARM64_FP16_VECTOR_ARITHMETIC
#define NEON_HALF_TO_FLOAT_FUNCTIONS(METHOD) \
template <> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE \
Packet8hf METHOD<Packet8hf>(const Packet8hf& x) { \
const Packet4f lo = METHOD<Packet4f>(vcvt_f32_f16(vget_low_f16(x))); \
const Packet4f hi = METHOD<Packet4f>(vcvt_f32_f16(vget_high_f16(x))); \
return vcombine_f16(vcvt_f16_f32(lo), vcvt_f16_f32(hi)); \
} \
\
template <> EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE \
Packet4hf METHOD<Packet4hf>(const Packet4hf& x) { \
return vcvt_f16_f32(METHOD<Packet4f>(vcvt_f32_f16(x))); \
}
NEON_HALF_TO_FLOAT_FUNCTIONS(perf)
NEON_HALF_TO_FLOAT_FUNCTIONS(pndtri)
#undef NEON_HALF_TO_FLOAT_FUNCTIONS
#endif
BF16_PACKET_FUNCTION(Packet4f, Packet4bf, perf)
BF16_PACKET_FUNCTION(Packet4f, Packet4bf, pndtri)
} // namespace internal
} // namespace Eigen
#endif // EIGEN_NEON_SPECIALFUNCTIONS_H