NEON: Round float to bfloat16 with RSHRN libeigen/eigen!3410
diff --git a/Eigen/src/Core/arch/NEON/PacketMath.h b/Eigen/src/Core/arch/NEON/PacketMath.h index ee0a78b..f2c2726 100644 --- a/Eigen/src/Core/arch/NEON/PacketMath.h +++ b/Eigen/src/Core/arch/NEON/PacketMath.h
@@ -4800,30 +4800,24 @@ } } // namespace detail +// Round to nearest even on the bfloat16 boundary. With lsb the lowest bit kept, the scalar +// formula is (p + 0x7fff + lsb) >> 16, and RSHRN computes (v + 0x8000) >> 16, so v = p - 1 + lsb, +// that is p minus one where lsb is clear. The caller guarantees that every NaN lane has a zero low +// half, which the FPU's NaN propagation does for bfloat16 operands: the rounding then leaves the +// NaN's high half alone. A NaN with a nonzero low half can instead round to +-inf or carry into a +// zero. +EIGEN_STRONG_INLINE Packet4bf F32ToBf16KeepNaN(const Packet4f& p) { + Packet4ui bits = vreinterpretq_u32_f32(p); + Packet4ui round_down = vceqq_u32(vandq_u32(bits, vdupq_n_u32(0x10000)), vdupq_n_u32(0)); + return vrshrn_n_u32(vaddq_u32(bits, round_down), 16); +} + +// F32ToBf16KeepNaN for any float: NaN lanes become the bfloat16 NaN 0x7fc0, as on AVX. The scalar +// conversion keeps the NaN's sign; doing that here costs an op on the critical path of every round +// trip, 25% of a cast loop and 6% of exp on a Cortex-A72, for a bit that carries no value. EIGEN_STRONG_INLINE Packet4bf F32ToBf16(const Packet4f& p) { - // See the scalar implementation in BFloat16.h for a comprehensible explanation - // of this fast rounding algorithm - Packet4ui input = Packet4ui(vreinterpretq_u32_f32(p)); - - // lsb = (input >> 16) & 1 - Packet4ui lsb = vandq_u32(vshrq_n_u32(input, 16), vdupq_n_u32(1)); - - // rounding_bias = 0x7fff + lsb - Packet4ui rounding_bias = vaddq_u32(lsb, vdupq_n_u32(0x7fff)); - - // input += rounding_bias - input = vaddq_u32(input, rounding_bias); - - // input = input >> 16 - input = vshrq_n_u32(input, 16); - - // Replace float-nans by bfloat16-nans, that is 0x7fc0 - const Packet4ui bf16_nan = vdupq_n_u32(0x7fc0); - const Packet4ui mask = vceqq_f32(p, p); - input = vbslq_u32(mask, input, bf16_nan); - - // output = static_cast<uint16_t>(input) - return vmovn_u32(input); + Packet4ui bits = vbslq_u32(vceqq_f32(p, p), vreinterpretq_u32_f32(p), vdupq_n_u32(0x7fc00000)); + return F32ToBf16KeepNaN(vreinterpretq_f32_u32(bits)); } // Discard the low 16 bits of each float. Only valid when every lane already holds an exact @@ -4974,42 +4968,42 @@ template <> EIGEN_STRONG_INLINE Packet4bf padd<Packet4bf>(const Packet4bf& a, const Packet4bf& b) { - return F32ToBf16(padd<Packet4f>(Bf16ToF32(a), Bf16ToF32(b))); + return F32ToBf16KeepNaN(padd<Packet4f>(Bf16ToF32(a), Bf16ToF32(b))); } template <> EIGEN_STRONG_INLINE Packet4bf psub<Packet4bf>(const Packet4bf& a, const Packet4bf& b) { - return F32ToBf16(psub<Packet4f>(Bf16ToF32(a), Bf16ToF32(b))); + return F32ToBf16KeepNaN(psub<Packet4f>(Bf16ToF32(a), Bf16ToF32(b))); } template <> EIGEN_STRONG_INLINE Packet4bf pmul<Packet4bf>(const Packet4bf& a, const Packet4bf& b) { - return F32ToBf16(pmul<Packet4f>(Bf16ToF32(a), Bf16ToF32(b))); + return F32ToBf16KeepNaN(pmul<Packet4f>(Bf16ToF32(a), Bf16ToF32(b))); } template <> EIGEN_STRONG_INLINE Packet4bf pmadd<Packet4bf>(const Packet4bf& a, const Packet4bf& b, const Packet4bf& c) { - return F32ToBf16(pmadd<Packet4f>(Bf16ToF32(a), Bf16ToF32(b), Bf16ToF32(c))); + return F32ToBf16KeepNaN(pmadd<Packet4f>(Bf16ToF32(a), Bf16ToF32(b), Bf16ToF32(c))); } template <> EIGEN_STRONG_INLINE Packet4bf pmsub<Packet4bf>(const Packet4bf& a, const Packet4bf& b, const Packet4bf& c) { - return F32ToBf16(pmsub<Packet4f>(Bf16ToF32(a), Bf16ToF32(b), Bf16ToF32(c))); + return F32ToBf16KeepNaN(pmsub<Packet4f>(Bf16ToF32(a), Bf16ToF32(b), Bf16ToF32(c))); } template <> EIGEN_STRONG_INLINE Packet4bf pnmadd<Packet4bf>(const Packet4bf& a, const Packet4bf& b, const Packet4bf& c) { - return F32ToBf16(pnmadd<Packet4f>(Bf16ToF32(a), Bf16ToF32(b), Bf16ToF32(c))); + return F32ToBf16KeepNaN(pnmadd<Packet4f>(Bf16ToF32(a), Bf16ToF32(b), Bf16ToF32(c))); } template <> EIGEN_STRONG_INLINE Packet4bf pnmsub<Packet4bf>(const Packet4bf& a, const Packet4bf& b, const Packet4bf& c) { - return F32ToBf16(pnmsub<Packet4f>(Bf16ToF32(a), Bf16ToF32(b), Bf16ToF32(c))); + return F32ToBf16KeepNaN(pnmsub<Packet4f>(Bf16ToF32(a), Bf16ToF32(b), Bf16ToF32(c))); } template <> EIGEN_STRONG_INLINE Packet4bf pdiv<Packet4bf>(const Packet4bf& a, const Packet4bf& b) { - return F32ToBf16(pdiv<Packet4f>(Bf16ToF32(a), Bf16ToF32(b))); + return F32ToBf16KeepNaN(pdiv<Packet4f>(Bf16ToF32(a), Bf16ToF32(b))); } template <> @@ -5058,7 +5052,7 @@ template <> EIGEN_STRONG_INLINE Packet4bf pabsdiff<Packet4bf>(const Packet4bf& a, const Packet4bf& b) { - return F32ToBf16(pabsdiff<Packet4f>(Bf16ToF32(a), Bf16ToF32(b))); + return F32ToBf16KeepNaN(pabsdiff<Packet4f>(Bf16ToF32(a), Bf16ToF32(b))); } template <>
diff --git a/Eigen/src/Core/arch/NEON/TypeCasting.h b/Eigen/src/Core/arch/NEON/TypeCasting.h index 5ea977e..5ab6bb7 100644 --- a/Eigen/src/Core/arch/NEON/TypeCasting.h +++ b/Eigen/src/Core/arch/NEON/TypeCasting.h
@@ -1740,6 +1740,21 @@ #endif // EIGEN_ARCH_ARM64 +template <> +struct type_casting_traits<bfloat16, float> : vectorized_type_casting_traits<bfloat16, float> {}; +template <> +struct type_casting_traits<float, bfloat16> : vectorized_type_casting_traits<float, bfloat16> {}; + +template <> +EIGEN_STRONG_INLINE Packet4f pcast<Packet4bf, Packet4f>(const Packet4bf& a) { + return Bf16ToF32(a); +} + +template <> +EIGEN_STRONG_INLINE Packet4bf pcast<Packet4f, Packet4bf>(const Packet4f& a) { + return F32ToBf16(a); +} + } // end namespace internal } // end namespace Eigen
diff --git a/benchmarks/Core/bench_cwise_math.cpp b/benchmarks/Core/bench_cwise_math.cpp index 2428af4..54a582a 100644 --- a/benchmarks/Core/bench_cwise_math.cpp +++ b/benchmarks/Core/bench_cwise_math.cpp
@@ -260,6 +260,9 @@ BENCH_CWISE_BINARY(Min, a.min(b), -100, 100) BENCH_CWISE_BINARY(Max, a.max(b), -100, 100) +BENCH_CWISE_BINARY(Add, a + b, -100, 100) +BENCH_CWISE_BINARY(Mul, (a * b), -100, 100) +BENCH_CWISE_BINARY(Div, a / b, 1, 100) // Power: array^scalar template <typename Scalar> @@ -356,6 +359,12 @@ ->ArgNames({"size", "mode"})->Name("Erf_float"); BENCHMARK(BM_Erf<bfloat16>)->ArgsProduct({{1024, 4096, 16384, 65536, 262144, 1048576}, {0, 1, 2}}) ->ArgNames({"size", "mode"})->Name("Erf_bfloat16"); +BENCHMARK(BM_Add<bfloat16>) CWISE_SIZES ->Name("Add_bfloat16"); +BENCHMARK(BM_Mul<bfloat16>) CWISE_SIZES ->Name("Mul_bfloat16"); +BENCHMARK(BM_Div<bfloat16>) CWISE_SIZES ->Name("Div_bfloat16"); +BENCHMARK(BM_Exp<bfloat16>) CWISE_SIZES ->Name("Exp_bfloat16"); +BENCHMARK(BM_Log<bfloat16>) CWISE_SIZES ->Name("Log_bfloat16"); +BENCHMARK(BM_Tanh<bfloat16>) CWISE_SIZES ->Name("Tanh_bfloat16"); BENCHMARK(BM_LogSpecialInputs<float>)->ArgsProduct({{4096, 65536}, {0, 1, 2}}) ->ArgNames({"size", "mode"})->Name("LogSpecialInputs_float"); BENCHMARK(BM_Log1pSpecialInputs<float>)->ArgsProduct({{4096, 65536}, {0, 1, 2}})
diff --git a/test/packetmath.cpp b/test/packetmath.cpp index bebd03a..d48ec43 100644 --- a/test/packetmath.cpp +++ b/test/packetmath.cpp
@@ -3121,6 +3121,66 @@ } #endif +#if defined(EIGEN_VECTORIZE_AVX) || defined(EIGEN_VECTORIZE_AVX512) || defined(EIGEN_VECTORIZE_NEON) || \ + defined(EIGEN_VECTORIZE_SME) +// The bfloat16 arithmetic packets round a float result to nearest even. Check them against that rounding applied to +// the float packet op, over every encoding of one operand and a fixed set for the other that puts sums and products +// on exact midpoints: 1 and 3 against [256, 512), 1 + 2^-7 and 2^-7 against values near 1, plus the range ends and +// the special values. NaN lanes compare as NaN, since a kernel may keep the operand's payload. +void packetmath_bfloat16_arithmetic_rounding() { + using Packet = internal::packet_traits<bfloat16>::type; + using FloatPacket = internal::packet_traits<float>::type; + constexpr int packet_size = internal::unpacket_traits<Packet>::size; + static_assert(internal::unpacket_traits<FloatPacket>::size == packet_size, + "the bfloat16 packet and the float packet must have the same size"); + const numext::uint16_t others[] = {0x0000, 0x8000, 0x3f80, 0x4040, 0x3f81, 0x3c00, 0x4380, 0xbf81, + 0x7f7f, 0x0080, 0x0001, 0x7f80, 0xff80, 0x7fc0, 0x7f81, 0x3f7f}; + constexpr int count = sizeof(others) / sizeof(others[0]); + EIGEN_ALIGN_MAX bfloat16 a[packet_size], b[packet_size], c[packet_size], out[packet_size]; + EIGEN_ALIGN_MAX float fa[packet_size], fb[packet_size], fc[packet_size], ref[packet_size]; + auto verify = [&](const Packet& result, const FloatPacket& expected) { + internal::pstore(out, result); + internal::pstore(ref, expected); + for (int i = 0; i < packet_size; ++i) { + numext::uint16_t got = numext::bit_cast<numext::uint16_t>(out[i]); + numext::uint16_t want = numext::bit_cast<numext::uint16_t>(static_cast<bfloat16>(ref[i])); + if ((want & 0x7fff) > 0x7f80) { + VERIFY((got & 0x7fff) > 0x7f80); + } else { + VERIFY_IS_EQUAL(got, want); + } + } + }; + for (int j = 0; j < count; ++j) { + for (int i = 0; i < packet_size; ++i) { + b[i] = numext::bit_cast<bfloat16>(others[j]); + c[i] = numext::bit_cast<bfloat16>(others[(j + 1) % count]); + fb[i] = static_cast<float>(b[i]); + fc[i] = static_cast<float>(c[i]); + } + Packet qb = internal::pload<Packet>(b), qc = internal::pload<Packet>(c); + FloatPacket pb = internal::pload<FloatPacket>(fb), pc = internal::pload<FloatPacket>(fc); + for (unsigned first = 0; first < 65536; first += packet_size) { + for (int i = 0; i < packet_size; ++i) { + a[i] = numext::bit_cast<bfloat16>(static_cast<numext::uint16_t>(first + i)); + fa[i] = static_cast<float>(a[i]); + } + Packet qa = internal::pload<Packet>(a); + FloatPacket pa = internal::pload<FloatPacket>(fa); + verify(internal::padd(qa, qb), internal::padd(pa, pb)); + verify(internal::psub(qa, qb), internal::psub(pa, pb)); + verify(internal::pmul(qa, qb), internal::pmul(pa, pb)); + verify(internal::pdiv(qa, qb), internal::pdiv(pa, pb)); + verify(internal::pabsdiff(qa, qb), internal::pabsdiff(pa, pb)); + verify(internal::pmadd(qa, qb, qc), internal::pmadd(pa, pb, pc)); + verify(internal::pmsub(qa, qb, qc), internal::pmsub(pa, pb, pc)); + verify(internal::pnmadd(qa, qb, qc), internal::pnmadd(pa, pb, pc)); + verify(internal::pnmsub(qa, qb, qc), internal::pnmsub(pa, pb, pc)); + } + } +} +#endif + namespace Eigen { namespace test { @@ -3200,6 +3260,11 @@ CALL_SUBTEST_15(packetmath_bfloat16_sign_bits_flushed()); #endif +#if defined(EIGEN_VECTORIZE_AVX) || defined(EIGEN_VECTORIZE_AVX512) || defined(EIGEN_VECTORIZE_NEON) || \ + defined(EIGEN_VECTORIZE_SME) + CALL_SUBTEST_15(packetmath_bfloat16_arithmetic_rounding()); +#endif + #if defined(EIGEN_VECTORIZE_RVV10) CALL_SUBTEST_1((packetmath_redux_infinities<float, internal::Packet1Xf>())); CALL_SUBTEST_1((packetmath_redux_infinities<float, internal::Packet2Xf>()));