new kernels for atan2

include/volk/volk_avx2_intrinsics.h

@@ -1,6 +1,7 @@
 /* -*- c++ -*- */
 /*
  * Copyright 2015 Free Software Foundation, Inc.
+ * Copyright 2023 Magnus Lundmark <[email protected]>
  *
  * This file is part of VOLK
  *
@@ -17,6 +18,20 @@
 #include "volk/volk_avx_intrinsics.h"
 #include <immintrin.h>
 
+static inline __m256 _mm256_real(const __m256 z1, const __m256 z2)
+{
+    const __m256i permute_mask = _mm256_set_epi32(7, 6, 3, 2, 5, 4, 1, 0);
+    __m256 r = _mm256_shuffle_ps(z1, z2, _MM_SHUFFLE(2, 0, 2, 0));
+    return _mm256_permutevar8x32_ps(r, permute_mask);
+}
+
+static inline __m256 _mm256_imag(const __m256 z1, const __m256 z2)
+{
+    const __m256i permute_mask = _mm256_set_epi32(7, 6, 3, 2, 5, 4, 1, 0);
+    __m256 i = _mm256_shuffle_ps(z1, z2, _MM_SHUFFLE(3, 1, 3, 1));
+    return _mm256_permutevar8x32_ps(i, permute_mask);
+}
+
 static inline __m256 _mm256_polar_sign_mask_avx2(__m128i fbits)
 {
     const __m128i zeros = _mm_set1_epi8(0x00);

include/volk/volk_common.h

@@ -170,7 +170,7 @@ static inline float log2f_non_ieee(float f)
 #define volk_log2to10factor (0x1.815182p1) // 3.01029995663981209120
 
 ////////////////////////////////////////////////////////////////////////
-// arctan(x)
+// arctan(x) polynomial expansion
 ////////////////////////////////////////////////////////////////////////
 static inline float volk_arctan_poly(const float x)
 {
@@ -198,19 +198,50 @@ static inline float volk_arctan_poly(const float x)
 
     return arctan;
 }
-
+////////////////////////////////////////////////////////////////////////
+// arctan(x)
+////////////////////////////////////////////////////////////////////////
 static inline float volk_arctan(const float x)
 {
     /*
      *  arctan(x) + arctan(1 / x) == sign(x) * pi / 2
      */
-    const float pi_over_2 = 0x1.921fb6p0f;
+    const float pi_2 = 0x1.921fb6p0f;
 
     if (fabs(x) < 1.f) {
         return volk_arctan_poly(x);
     } else {
-        return copysignf(pi_over_2, x) - volk_arctan_poly(1.f / x);
+        return copysignf(pi_2, x) - volk_arctan_poly(1.f / x);
+    }
+}
+////////////////////////////////////////////////////////////////////////
+// arctan2(y, x)
+////////////////////////////////////////////////////////////////////////
+static inline float volk_atan2(const float y, const float x)
+{
+    /*
+     *                /  arctan(y / x)         if x > 0
+     *                |  arctan(y / x)  + PI   if x < 0 and y >= 0
+     * atan2(y, x) =  |  arctan(y / x)  - PI   if x < 0 and y <  0
+     *                |  sign(y) * PI / 2      if x = 0
+     *                \  undefined             if x = 0 and y = 0
+     * atan2f(0.f,  0.f) shall return  0.f
+     * atan2f(0.f, -0.f) shall return -0.f
+     */
+    const float pi = 0x1.921fb6p1f;
+    const float pi_2 = 0x1.921fb6p0f;
+
+    if (fabs(x) == 0.f) {
+        return (fabs(y) == 0.f) ? copysignf(0.f, y) : copysignf(pi_2, y);
+    }
+    const int swap = fabs(x) < fabs(y);
+    const float input = swap ? (x / y) : (y / x);
+    float result = volk_arctan_poly(input);
+    result = swap ? (input >= 0.f ? pi_2 : -pi_2) - result : result;
+    if (x < 0.f) {
+        result += copysignf(pi, y);
     }
+    return result;
 }
 
 #endif /*INCLUDED_LIBVOLK_COMMON_H*/