| #ifndef EIGEN_PACKET_MATH_GOOGLE_AVX_H |
| #define EIGEN_PACKET_MATH_GOOGLE_AVX_H |
| |
| #include <stdio.h> |
| #include <x86intrin.h> |
| |
| #include "Eigen/src/Core/util/Macros.h" |
| |
| namespace Eigen { |
| |
| namespace internal { |
| |
| typedef __m256 Packet8f; |
| |
| #define BFLOAT_MASK 0xFFFF0000 |
| #define BFLOAT_ROUND_MASK 0xFFFF8000 |
| #define BFLOAT_EXP_MASK 0xFF800000 |
| #define BFLOAT_HALF_MASK 0x00008000 |
| #define BFLOAT_SIGN_BIT 0x80000000 |
| |
| #if defined(CUSTOM_NUMERICS) |
| #if !defined(CUSTOM_NUMERICS_MANT_BITS) |
| #define CUSTOM_NUMERICS_MANT_BITS 7 |
| #endif // CUSTOM_NUMERICS_MANT_BITS |
| #if !defined(CUSTOM_NUMERICS_EXP_BITS) |
| #define CUSTOM_NUMERICS_EXP_BITS 8 |
| #endif // CUSTOM_NUMERICS_EXP_BITS |
| #if !defined(CUSTOM_NUMERICS_EXP_BIAS) |
| #define CUSTOM_NUMERICS_EXP_BIAS ((1 << (CUSTOM_NUMERICS_EXP_BITS - 1)) - 1) |
| #endif // CUSTOM_NUMERICS_EXP_BIAS |
| #endif // CUSTOM_NUMERICS |
| |
| EIGEN_STRONG_INLINE Packet8f BFConvertTrunc(const Packet8f& value) { |
| #if defined(HALF_PRECISION_BF16) |
| const __m256 bfloat_mask = _mm256_castsi256_ps( |
| _mm256_set1_epi32(BFLOAT_MASK)); |
| __m256 result = _mm256_and_ps(bfloat_mask, value); |
| return result; |
| #elif defined(CUSTOM_NUMERICS) && (CUSTOM_NUMERICS_EXP_BITS == 8) && \ |
| (CUSTOM_NUMERICS_EXP_BIAS == 127) |
| const int mask_scalar_ = 0xFFFFFFFF << (23 - CUSTOM_NUMERICS_MANT_BITS); |
| const __m256 bfloat_mask = _mm256_castsi256_ps( |
| _mm256_set1_epi32(mask_scalar_)); |
| __m256 result = _mm256_and_ps(bfloat_mask, value); |
| return result; |
| #elif defined(CUSTOM_NUMERICS) |
| const __m256 sign_bit_simd_ = _mm256_castsi256_ps( |
| _mm256_set1_epi32(BFLOAT_SIGN_BIT)); |
| const int mask_scalar_ = 0xFFFFFFFF << (23 - CUSTOM_NUMERICS_MANT_BITS); |
| const int abs_mask_scalar_ = mask_scalar_ & 0x7FFFFFFF; |
| const __m256 abs_mask_simd_ = |
| _mm256_castsi256_ps(_mm256_set1_epi32(abs_mask_scalar_)); |
| const __m256 zero_simd_ = _mm256_castsi256_ps(_mm256_set1_epi32(0)); |
| const int min_float_i = (128 - CUSTOM_NUMERICS_EXP_BIAS) << 23; |
| const int max_float_i = |
| (((127 + (1 << CUSTOM_NUMERICS_EXP_BITS) - 2 - CUSTOM_NUMERICS_EXP_BIAS) |
| << 23) | ((1 << 23) - 1)) & mask_scalar_; |
| const __m256 min_simd_ = |
| _mm256_castsi256_ps(_mm256_set1_epi32(min_float_i)); |
| const __m256 max_simd_ = |
| _mm256_castsi256_ps(_mm256_set1_epi32(max_float_i)); |
| // extract sign bit and get the abs value |
| __m256 sign = _mm256_and_ps(sign_bit_simd_, value); |
| __m256 abs_value = _mm256_and_ps(abs_mask_simd_, value); |
| // mask the ones equal to 0 for SIMD construction |
| __m256 mask = _mm256_cmp_ps(abs_value, zero_simd_, 0); |
| // check against exponent range |
| __m256 result = _mm256_max_ps(abs_value, min_simd_); |
| result = _mm256_min_ps(result, max_simd_); |
| // apply the sign bit |
| result = _mm256_or_ps(result, sign); |
| // blend back to SIMD float with proper 0 value |
| result = _mm256_blendv_ps(result, zero_simd_, mask); |
| return result; |
| #else |
| return value; |
| #endif // HALF_PRECISION_BF16 |
| } |
| |
| EIGEN_STRONG_INLINE Packet8f BFConvertRound(const Packet8f& value) { |
| #if defined(HALF_PRECISION_BF16) |
| const __m256 bfloat_mask = _mm256_castsi256_ps( |
| _mm256_set1_epi32(BFLOAT_MASK)); |
| const __m256 bfloat_round_mask = _mm256_castsi256_ps( |
| _mm256_set1_epi32(BFLOAT_ROUND_MASK)); |
| const __m256 bfloat_exp_mask = _mm256_castsi256_ps( |
| _mm256_set1_epi32(BFLOAT_EXP_MASK)); |
| const __m256 bfloat_half_mask = _mm256_castsi256_ps( |
| _mm256_set1_epi32(BFLOAT_HALF_MASK)); |
| __m256 result = _mm256_and_ps(bfloat_round_mask, value); |
| __m256 one_scaled = _mm256_and_ps(bfloat_exp_mask, value); |
| __m256 one_plus_half_scaled = _mm256_or_ps(one_scaled, |
| bfloat_half_mask); |
| // subtract 1 (scaled) |
| result = _mm256_sub_ps(result, one_scaled); |
| // add 1 + half the lsb (scaled) |
| result = _mm256_add_ps(result, one_plus_half_scaled); |
| // truncate after adding half lsb |
| result = _mm256_and_ps(bfloat_mask, result); |
| return result; |
| #elif defined(CUSTOM_NUMERICS) && (CUSTOM_NUMERICS_EXP_BITS == 8) && \ |
| (CUSTOM_NUMERICS_EXP_BIAS == 127) |
| const int mask_scalar_ = 0xFFFFFFFF << (23 - CUSTOM_NUMERICS_MANT_BITS); |
| const int mask_scalar_r_ = 0xFFFFFFFF << (22 - CUSTOM_NUMERICS_MANT_BITS); |
| const int mask_scalar_half_ = 0x1 << (22 - CUSTOM_NUMERICS_MANT_BITS); |
| const __m256 bfloat_mask = _mm256_castsi256_ps( |
| _mm256_set1_epi32(mask_scalar_)); |
| const __m256 bfloat_round_mask = _mm256_castsi256_ps( |
| _mm256_set1_epi32(mask_scalar_r_)); |
| const __m256 bfloat_exp_mask = _mm256_castsi256_ps( |
| _mm256_set1_epi32(BFLOAT_EXP_MASK)); |
| const __m256 bfloat_half_mask = _mm256_castsi256_ps( |
| _mm256_set1_epi32(mask_scalar_half_)); |
| __m256 result = _mm256_and_ps(bfloat_round_mask, value); |
| __m256 one_scaled = _mm256_and_ps(bfloat_exp_mask, value); |
| __m256 one_plus_half_scaled = _mm256_or_ps(one_scaled, |
| bfloat_half_mask); |
| // subtract 1 (scaled) |
| result = _mm256_sub_ps(result, one_scaled); |
| // add 1 + half the lsb (scaled) |
| result = _mm256_add_ps(result, one_plus_half_scaled); |
| // truncate after adding half lsb |
| result = _mm256_and_ps(bfloat_mask, result); |
| return result; |
| #elif defined(CUSTOM_NUMERICS) |
| const __m256 sign_bit_simd_ = _mm256_castsi256_ps( |
| _mm256_set1_epi32(BFLOAT_SIGN_BIT)); |
| const int mask_scalar_ = 0xFFFFFFFF << (23 - CUSTOM_NUMERICS_MANT_BITS); |
| const int abs_mask_scalar_ = mask_scalar_ & 0x7FFFFFFF; |
| const __m256 abs_mask_simd_ = |
| _mm256_castsi256_ps(_mm256_set1_epi32(abs_mask_scalar_)); |
| const int half_bit_scalar_ = |
| (mask_scalar_ ^ (mask_scalar_ >> 1)) & 0x7FFFFFFF; |
| const int mask_plus_half_scalar_ = mask_scalar_ | half_bit_scalar_; |
| const __m256 abs_mask_plus_half_simd_ = |
| _mm256_castsi256_ps(_mm256_set1_epi32(mask_plus_half_scalar_)); |
| const __m256 half_simd_ = |
| _mm256_castsi256_ps(_mm256_set1_epi32(half_bit_scalar_)); |
| const __m256 exp_mask_simd_ = |
| _mm256_castsi256_ps(_mm256_set1_epi32(0x7F800000)); |
| const __m256 zero_simd_ = _mm256_castsi256_ps(_mm256_set1_epi32(0)); |
| const int min_float_i = (128 - CUSTOM_NUMERICS_EXP_BIAS) << 23; |
| const int max_float_i = |
| (((127 + (1 << CUSTOM_NUMERICS_EXP_BITS) - 2 - CUSTOM_NUMERICS_EXP_BIAS) |
| << 23) | ((1 << 23) - 1)) & mask_scalar_; |
| const __m256 min_simd_ = |
| _mm256_castsi256_ps(_mm256_set1_epi32(min_float_i)); |
| const __m256 max_simd_ = |
| _mm256_castsi256_ps(_mm256_set1_epi32(max_float_i)); |
| // extract sign bit and get the abs value |
| __m256 sign = _mm256_and_ps(sign_bit_simd_, value); |
| __m256 abs_value = _mm256_and_ps(abs_mask_plus_half_simd_, value); |
| __m256 one_scaled = _mm256_and_ps(exp_mask_simd_, value); |
| __m256 one_plus_half_scaled = _mm256_or_ps(one_scaled, half_simd_); |
| // subtract 1 (scaled) |
| abs_value = _mm256_sub_ps(abs_value, one_scaled); |
| // add 1 + half the lsb (scaled) |
| abs_value = _mm256_add_ps(abs_value, one_plus_half_scaled); |
| // truncate after adding half lsb |
| abs_value = _mm256_and_ps(abs_mask_simd_, abs_value); |
| // mask the ones equal to 0 for SIMD construction |
| __m256 mask = _mm256_cmp_ps(abs_value, zero_simd_, 0); |
| // check against exponent range |
| __m256 result = _mm256_max_ps(abs_value, min_simd_); |
| result = _mm256_min_ps(result, max_simd_); |
| // apply the sign bit |
| result = _mm256_or_ps(result, sign); |
| // blend back to SIMD float with proper 0 value |
| result = _mm256_blendv_ps(result, zero_simd_, mask); |
| return result; |
| #else |
| return value; |
| #endif // HALF_PRECISION_BF16 |
| } |
| |
| EIGEN_STRONG_INLINE Packet8f BFConvert(const Packet8f& value) { |
| #if defined(HALF_PRECISION_ROUND) || defined(CUSTOM_NUMERICS_ROUND) |
| return BFConvertRound(value); |
| #else |
| return BFConvertTrunc(value); |
| #endif // HALF_PRECISION_ROUND || defined(CUSTOM_NUMERICS_ROUND) |
| } |
| |
| EIGEN_STRONG_INLINE Packet8f pmul_half_bf16( |
| const Packet8f& a, const Packet8f& b) { |
| __m256 a_bf16 = BFConvert(a); |
| __m256 b_bf16 = BFConvert(b); |
| #if defined(HALF_PRECISION_OUT16) || defined(CUSTOM_NUMERICS_OUT) |
| return BFConvert(_mm256_mul_ps(a_bf16, b_bf16)); |
| #else |
| return _mm256_mul_ps(a_bf16, b_bf16); |
| #endif // HALF_PRECISION_OUT16 || defined(CUSTOM_NUMERICS_OUT) |
| } |
| |
| EIGEN_STRONG_INLINE Packet8f pmul_half_fp16( |
| const Packet8f& a, const Packet8f& b) { |
| #if defined(HALF_PRECISION_FP16) |
| #if defined(HALF_PRECISION_ROUND) || defined(CUSTOM_NUMERICS_ROUND) |
| // default round to nearest even is _MM_FROUND_TO_NEAREST_INT = 0x0 |
| __m128i a_fp16 = _mm256_cvtps_ph(a, _MM_FROUND_NO_EXC); |
| __m128i b_fp16 = _mm256_cvtps_ph(b, _MM_FROUND_NO_EXC); |
| #else |
| __m128i a_fp16 = _mm256_cvtps_ph(a, _MM_FROUND_TO_ZERO |_MM_FROUND_NO_EXC); |
| __m128i b_fp16 = _mm256_cvtps_ph(b, _MM_FROUND_TO_ZERO |_MM_FROUND_NO_EXC); |
| #endif // HALF_PRECISION_ROUND || defined(CUSTOM_NUMERICS_ROUND) |
| __m256 tmp = _mm256_mul_ps(_mm256_cvtph_ps(a_fp16), _mm256_cvtph_ps(b_fp16)); |
| #if defined(HALF_PRECISION_OUT16) || defined(CUSTOM_NUMERICS_OUT) |
| #if defined(HALF_PRECISION_ROUND) || defined(CUSTOM_NUMERICS_ROUND) |
| __m128i tmp_fp16 = _mm256_cvtps_ph(tmp, _MM_FROUND_NO_EXC); |
| #else |
| __m128i tmp_fp16 = _mm256_cvtps_ph(tmp, |
| _MM_FROUND_TO_ZERO |_MM_FROUND_NO_EXC); |
| #endif // HALF_PRECISION_ROUND || defined(CUSTOM_NUMERICS_ROUND) |
| return _mm256_cvtph_ps(tmp_fp16); |
| #else |
| return tmp; |
| #endif // HALF_PRECISION_OUT16 || defined(CUSTOM_NUMERICS_OUT) |
| #else |
| return _mm256_mul_ps(a, b); |
| #endif // HALF_PRECISION_FP16 |
| } |
| |
| EIGEN_STRONG_INLINE Packet8f pmul_custom(const Packet8f& a, const Packet8f& b) { |
| #if defined(HALF_PRECISION_BF16) || defined(CUSTOM_NUMERICS) |
| return pmul_half_bf16(a, b); |
| #else |
| return pmul_half_fp16(a, b); |
| #endif // HALF_PRECISION_BF16 || defined(CUSTOM_NUMERICS) |
| } |
| |
| EIGEN_STRONG_INLINE Packet8f pmadd_half_bf16( |
| const Packet8f& a, const Packet8f& b, const Packet8f& c) { |
| __m256 a_bf16 = BFConvert(a); |
| __m256 b_bf16 = BFConvert(b); |
| #if defined(HALF_PRECISION_OUT16) || defined(CUSTOM_NUMERICS_OUT) |
| return BFConvert(_mm256_fmadd_ps(a_bf16, b_bf16, c)); |
| #else |
| return _mm256_fmadd_ps(a_bf16, b_bf16, c); |
| #endif // HALF_PRECISION_OUT16 || defined(CUSTOM_NUMERICS_OUT) |
| } |
| |
| EIGEN_STRONG_INLINE Packet8f pmadd_half_fp16( |
| const Packet8f& a, const Packet8f& b, const Packet8f& c) { |
| #if defined(HALF_PRECISION_FP16) |
| #if defined(HALF_PRECISION_ROUND) || defined(CUSTOM_NUMERICS_ROUND) |
| // default round to nearest even is _MM_FROUND_TO_NEAREST_INT = 0x0 |
| __m128i a_fp16 = _mm256_cvtps_ph(a, _MM_FROUND_NO_EXC); |
| __m128i b_fp16 = _mm256_cvtps_ph(b, _MM_FROUND_NO_EXC); |
| #else |
| __m128i a_fp16 = _mm256_cvtps_ph(a, _MM_FROUND_TO_ZERO |_MM_FROUND_NO_EXC); |
| __m128i b_fp16 = _mm256_cvtps_ph(b, _MM_FROUND_TO_ZERO |_MM_FROUND_NO_EXC); |
| #endif // HALF_PRECISION_ROUND || defined(CUSTOM_NUMERICS_ROUND) |
| __m256 tmp = _mm256_fmadd_ps(_mm256_cvtph_ps(a_fp16), |
| _mm256_cvtph_ps(b_fp16), |
| c); |
| #if defined(HALF_PRECISION_OUT16) || defined(CUSTOM_NUMERICS_OUT) |
| #if defined(HALF_PRECISION_ROUND) || defined(CUSTOM_NUMERICS_ROUND) |
| __m128i tmp_fp16 = _mm256_cvtps_ph(tmp, _MM_FROUND_NO_EXC); |
| #else |
| __m128i tmp_fp16 = _mm256_cvtps_ph(tmp, |
| _MM_FROUND_TO_ZERO |_MM_FROUND_NO_EXC); |
| #endif // HALF_PRECISION_ROUND || defined(CUSTOM_NUMERICS_ROUND) |
| return _mm256_cvtph_ps(tmp_fp16); |
| #else |
| return tmp; |
| #endif // HALF_PRECISION_OUT16 || defined(CUSTOM_NUMERICS_OUT) |
| #else |
| return _mm256_fmadd_ps(a, b, c); |
| #endif // HALF_PRECISION_FP16 |
| } |
| |
| EIGEN_STRONG_INLINE Packet8f pmadd_custom( |
| const Packet8f& a, const Packet8f& b, const Packet8f& c) { |
| #if defined(HALF_PRECISION_BF16) || defined(CUSTOM_NUMERICS) |
| return pmadd_half_bf16(a, b, c); |
| #else |
| return pmadd_half_fp16(a, b, c); |
| #endif // HALF_PRECISION_BF16 || defined(CUSTOM_NUMERICS) |
| } |
| |
| } // end namespace internal |
| |
| } // end namespace Eigen |
| |
| #endif // EIGEN_PACKET_MATH_AVX_H |
