SpecialFunctions: Preserve subnormal erf results on ARMv7 NEON

libeigen/eigen!3040

diff --git a/benchmarks/Core/bench_cwise_math.cpp b/benchmarks/Core/bench_cwise_math.cpp
index 34b0ba6..ea28234 100644
--- a/benchmarks/Core/bench_cwise_math.cpp
+++ b/benchmarks/Core/bench_cwise_math.cpp
@@ -58,7 +58,31 @@
 BENCH_CWISE_UNARY(Acosh, a.acosh(), 1.01, 10)
 BENCH_CWISE_UNARY(Atanh, a.atanh(), -0.99, 0.99)
 BENCH_CWISE_UNARY(Log10, a.log10(), 0.01, 100)
-BENCH_CWISE_UNARY(Erf, Eigen::erf(a), -4, 4)
+
+// mode: 0 = ordinary inputs, 1 = subnormal inputs, 2 = one subnormal per 64 coefficients.
+template <typename Scalar>
+static void BM_Erf(benchmark::State& state) {
+  using Bits = typename numext::get_integer_by_size<sizeof(Scalar)>::unsigned_type;
+  constexpr Bits sign = Bits(1) << (8 * sizeof(Scalar) - 1);
+  constexpr Bits minNormal = Bits(1) << (std::numeric_limits<Scalar>::digits - 1);
+  const Index n = state.range(0);
+  const int mode = int(state.range(1));
+  Array<Scalar, Dynamic, 1> a(n), b(n);
+  for (Index i = 0; i < n; ++i) {
+    a(i) = Scalar(double(i % 257 - 128) / 32.0);
+    if (mode == 1 || (mode == 2 && i % 64 == 0)) {
+      const Bits magnitude = Bits(1 + (Bits(i) * 37) % (minNormal - 1));
+      a(i) = numext::bit_cast<Scalar>(Bits(magnitude | (i % 2 ? sign : Bits(0))));
+    }
+  }
+  for (auto _ : state) {
+    benchmark::DoNotOptimize(a.data());
+    b = Eigen::erf(a);
+    benchmark::DoNotOptimize(b.data());
+    benchmark::ClobberMemory();
+  }
+  state.SetBytesProcessed(state.iterations() * n * sizeof(Scalar) * 2);
+}
 
 // Simple operations (should be very fast / memory-bound)
 BENCH_CWISE_UNARY(Abs, a.abs(), -100, 100)
@@ -182,7 +206,10 @@
 BENCHMARK(BM_Acosh<float>) CWISE_SIZES ->Name("Acosh_float");
 BENCHMARK(BM_Atanh<float>) CWISE_SIZES ->Name("Atanh_float");
 BENCHMARK(BM_Log10<float>) CWISE_SIZES ->Name("Log10_float");
-BENCHMARK(BM_Erf<float>) CWISE_SIZES ->Name("Erf_float");
+BENCHMARK(BM_Erf<float>)->ArgsProduct({{1024, 4096, 16384, 65536, 262144, 1048576}, {0, 1, 2}})
+    ->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_Abs<float>) CWISE_SIZES ->Name("Abs_float");
 BENCHMARK(BM_Square<float>) CWISE_SIZES ->Name("Square_float");
 BENCHMARK(BM_Cube<float>) CWISE_SIZES ->Name("Cube_float");
@@ -219,7 +246,8 @@
 BENCHMARK(BM_Acosh<double>) CWISE_SIZES ->Name("Acosh_double");
 BENCHMARK(BM_Atanh<double>) CWISE_SIZES ->Name("Atanh_double");
 BENCHMARK(BM_Log10<double>) CWISE_SIZES ->Name("Log10_double");
-BENCHMARK(BM_Erf<double>) CWISE_SIZES ->Name("Erf_double");
+BENCHMARK(BM_Erf<double>)->ArgsProduct({{1024, 4096, 16384, 65536, 262144, 1048576}, {0, 1, 2}})
+    ->ArgNames({"size", "mode"})->Name("Erf_double");
 BENCHMARK(BM_Abs<double>) CWISE_SIZES ->Name("Abs_double");
 BENCHMARK(BM_Square<double>) CWISE_SIZES ->Name("Square_double");
 BENCHMARK(BM_Cube<double>) CWISE_SIZES ->Name("Cube_double");
diff --git a/contrib/Eigen/src/SpecialFunctions/arch/NEON/SpecialFunctions.h b/contrib/Eigen/src/SpecialFunctions/arch/NEON/SpecialFunctions.h
index 88d1d6f..3b27b9c 100644
--- a/contrib/Eigen/src/SpecialFunctions/arch/NEON/SpecialFunctions.h
+++ b/contrib/Eigen/src/SpecialFunctions/arch/NEON/SpecialFunctions.h
@@ -7,6 +7,42 @@
 namespace Eigen {
 namespace internal {
 
+#if EIGEN_ARCH_ARM
+// For x = n * 2^-149, n < 2^23, erf(x) = (2/sqrt(pi))*x to float precision.
+// round((2/sqrt(pi))*2^31) = 2423175810; the coefficient error contributes < 0.000319 ULP.
+// Integer rounding avoids ARMv7 NEON's input/output flushing. The result bits also cover the first normal binade.
+EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Packet2i perf_subnormal_significands(const Packet2i& n) {
+  return vreinterpret_s32_u32(vrshrn_n_u64(vmull_u32(vreinterpret_u32_s32(n), vdup_n_u32(2423175810u)), 31));
+}
+
+EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Packet4i perf_subnormal_significands(const Packet4i& n) {
+  return vcombine_s32(perf_subnormal_significands(Packet2i(vget_low_s32(n))),
+                      perf_subnormal_significands(Packet2i(vget_high_s32(n))));
+}
+
+template <typename Packet>
+EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Packet perf_neon(const Packet& x) {
+  using PacketI = typename unpacket_traits<Packet>::integer_packet;
+  const Packet result = generic_fast_erf<float>::run(x);
+  const PacketI bits = preinterpret<PacketI>(x);
+  const PacketI magnitude = pand(bits, pset1<PacketI>(0x7fffffff));
+  const PacketI sign = pxor(bits, magnitude);
+  const PacketI subnormal = pcmp_lt(magnitude, pset1<PacketI>(0x00800000));
+  const PacketI recovered = por(perf_subnormal_significands(magnitude), sign);
+  return pselect(preinterpret<Packet>(subnormal), preinterpret<Packet>(recovered), result);
+}
+
+template <>
+EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Packet2f perf<Packet2f>(const Packet2f& x) {
+  return perf_neon(x);
+}
+
+template <>
+EIGEN_DEVICE_FUNC EIGEN_STRONG_INLINE Packet4f perf<Packet4f>(const Packet4f& x) {
+  return perf_neon(x);
+}
+#endif
+
 #if EIGEN_ARCH_ARM64 && EIGEN_HAS_ARM64_FP16
 
 #define NEON_HALF_TO_FLOAT_FUNCTIONS(METHOD)                                              \
diff --git a/contrib/test/special_packetmath.cpp b/contrib/test/special_packetmath.cpp
index 9b500af..ba20772 100644
--- a/contrib/test/special_packetmath.cpp
+++ b/contrib/test/special_packetmath.cpp
@@ -11,10 +11,102 @@
 
 #include <limits>
 #include "packetmath_test_shared.h"
+#include "fp_control.h"
 #include "../Eigen/SpecialFunctions"
 
 using internal::unpacket_traits;
 
+#if EIGEN_ARCH_ARM && defined(EIGEN_VECTORIZE_NEON)
+template <typename Scalar, typename Packet>
+void packet_erf_subnormals() {
+  using Bits = typename numext::get_integer_by_size<sizeof(Scalar)>::unsigned_type;
+  constexpr int PacketSize = unpacket_traits<Packet>::size;
+  constexpr Bits Sign = Bits(1) << (8 * sizeof(Scalar) - 1);
+  constexpr Bits MinNormal = Bits(1) << (std::numeric_limits<Scalar>::digits - 1);
+  const Bits samples[] = {Bits(1), Bits(2), Bits(3), Bits(MinNormal / 2), Bits(MinNormal - 1)};
+  // In this interval the cubic term of erf is negligible even relative to one subnormal ULP.
+  const double slope = 2.0 / std::sqrt(std::acos(-1.0));
+  for (Bits n : samples) {
+    for (Bits sign : {Bits(0), Sign}) {
+      for (int lane = 0; lane < PacketSize; ++lane) {
+        Scalar input[PacketSize], output[PacketSize];
+        for (int i = 0; i < PacketSize; ++i) input[i] = Scalar(0.5);
+        input[lane] = numext::bit_cast<Scalar>(Bits(sign | n));
+        internal::pstoreu(output, internal::perf(internal::ploadu<Packet>(input)));
+        const Bits actual = numext::bit_cast<Bits>(output[lane]);
+        const Bits expected = Bits(std::floor(double(n) * slope + 0.5));
+        VERIFY_IS_EQUAL(Bits(actual & Sign), sign);
+        const int error = int(actual & (Sign - 1)) - int(expected);
+        VERIFY(std::abs(error) <= (sizeof(Scalar) == sizeof(float) ? 1 : 0));
+        if (n == 1) VERIFY_IS_EQUAL(actual, Bits(sign | n));
+        for (int i = 0; i < PacketSize; ++i) {
+          if (i != lane) VERIFY_IS_APPROX(output[i], Scalar(std::erf(0.5)));
+        }
+      }
+    }
+  }
+
+  // Recovering one lane must not change zeros, infinities, or NaNs in its neighbors.
+  const Scalar neighbors[] = {Scalar(0), numext::bit_cast<Scalar>(Sign), std::numeric_limits<Scalar>::infinity(),
+                              std::numeric_limits<Scalar>::quiet_NaN()};
+  for (Scalar neighbor : neighbors) {
+    Scalar values[PacketSize], before[PacketSize], after[PacketSize];
+    for (int i = 0; i < PacketSize; ++i) values[i] = neighbor;
+    internal::pstoreu(before, internal::perf(internal::ploadu<Packet>(values)));
+    if ((numext::bit_cast<Bits>(neighbor) & (Sign - 1)) == 0) {
+      for (int i = 0; i < PacketSize; ++i)
+        VERIFY_IS_EQUAL(numext::bit_cast<Bits>(before[i]), numext::bit_cast<Bits>(neighbor));
+    }
+    values[0] = numext::bit_cast<Scalar>(Bits(1));
+    internal::pstoreu(after, internal::perf(internal::ploadu<Packet>(values)));
+    VERIFY_IS_EQUAL(numext::bit_cast<Bits>(after[0]), Bits(1));
+    for (int i = 1; i < PacketSize; ++i)
+      VERIFY_IS_EQUAL(numext::bit_cast<Bits>(after[i]), numext::bit_cast<Bits>(before[i]));
+  }
+
+  // Check the public evaluator with packet bodies and an ordinary scalar tail, also while scalar arithmetic flushes.
+  Array<Scalar, Dynamic, 1> input(3 * internal::packet_traits<Scalar>::size + 1), output(input.size());
+  for (Index i = 0; i + 1 < input.size(); ++i)
+    input(i) = numext::bit_cast<Scalar>(Bits(Bits(1) | (i % 2 ? Sign : Bits(0))));
+  input(input.size() - 1) = Scalar(0.5);
+  output = input.erf();
+  for (Index i = 0; i + 1 < input.size(); ++i)
+    VERIFY_IS_EQUAL(numext::bit_cast<Bits>(output(i)), numext::bit_cast<Bits>(input(i)));
+  VERIFY_IS_APPROX(output(output.size() - 1), Scalar(std::erf(0.5)));
+}
+
+void neon_erf_subnormals_float() {
+  packet_erf_subnormals<float, internal::Packet2f>();
+  packet_erf_subnormals<float, internal::Packet4f>();
+
+  // Rounded values straddling the subnormal/normal output boundary and the largest subnormal input.
+  const numext::uint32_t inputs[] = {0x00716fe1u, 0x00716fe2u, 0x00716fe3u, 0x007fffffu};
+  const numext::uint32_t expected[] = {0x007fffffu, 0x00800000u, 0x00800001u, 0x00906eb9u};
+  float input[4], output[4];
+  for (int i = 0; i < 4; ++i) input[i] = numext::bit_cast<float>(inputs[i]);
+  internal::pstoreu(output, internal::perf(internal::ploadu<internal::Packet4f>(input)));
+  for (int i = 0; i < 4; ++i) {
+    const int error = int(numext::bit_cast<numext::uint32_t>(output[i])) - int(expected[i]);
+    VERIFY(std::abs(error) <= 1);
+  }
+}
+
+void neon_erf_subnormals_bfloat16() {
+  packet_erf_subnormals<bfloat16, internal::Packet4bf>();
+  // Exhaustive bfloat16 subnormals, with both signs, checked without floating-point subnormal conversions.
+  for (unsigned int n = 1; n < 128; ++n) {
+    const numext::uint16_t expectedBits =
+        numext::uint16_t(std::floor(double(n) * (2.0 / std::sqrt(std::acos(-1.0))) + 0.5));
+    bfloat16 values[4], result[4];
+    for (int i = 0; i < 4; ++i) values[i] = numext::bit_cast<bfloat16>(numext::uint16_t(n | (i % 2 ? 0x8000 : 0)));
+    internal::pstoreu(result, internal::perf(internal::ploadu<internal::Packet4bf>(values)));
+    for (int i = 0; i < 4; ++i)
+      VERIFY_IS_EQUAL(numext::bit_cast<numext::uint16_t>(result[i]),
+                      numext::uint16_t(expectedBits | (i % 2 ? 0x8000 : 0)));
+  }
+}
+#endif
+
 #if EIGEN_ARCH_ARM
 // Note: 32-bit arm always flushes subnormals to zero.
 #define MAYBE_FLUSH(op)                                                      \
@@ -213,6 +305,15 @@
 }  // namespace Eigen
 
 EIGEN_DECLARE_TEST(special_packetmath) {
+#if EIGEN_ARCH_ARM && defined(EIGEN_VECTORIZE_NEON)
+  CALL_SUBTEST_1(neon_erf_subnormals_float());
+  CALL_SUBTEST_4(neon_erf_subnormals_bfloat16());
+  {
+    const Eigen::ScopedFlushToZero flush_to_zero;
+    CALL_SUBTEST_1(neon_erf_subnormals_float());
+    CALL_SUBTEST_4(neon_erf_subnormals_bfloat16());
+  }
+#endif
   g_first_pass = true;
   for (int i = 0; i < g_repeat; i++) {
     CALL_SUBTEST_1(test::runner<float>::run());