Tweak atan2

Eigen/src/Core/GenericPacketMath.h

@@ -1219,6 +1219,56 @@ template <typename Packet>
 EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE constexpr Packet
 psignbit(const Packet& a) { return psignbit_impl<Packet>::run(a); }
 
+/** \internal \returns the 2-argument arc tangent of \a y and \a x (coeff-wise) */
+template <typename Packet, std::enable_if_t<is_scalar<Packet>::value, int> = 0>
+EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE Packet patan2(const Packet& y, const Packet& x) {
+  return numext::atan2(y, x);
+}
+
+/** \internal \returns the 2-argument arc tangent of \a y and \a x (coeff-wise) */
+template <typename Packet, std::enable_if_t<!is_scalar<Packet>::value, int> = 0>
+EIGEN_DEVICE_FUNC EIGEN_ALWAYS_INLINE Packet patan2(const Packet& y, const Packet& x) {
+  typedef typename internal::unpacket_traits<Packet>::type Scalar;
+
+  // See https://en.cppreference.com/w/cpp/numeric/math/atan2
+  // for how corner cases are supposed to be handled according to the
+  // IEEE floating-point standard (IEC 60559).
+
+  const Packet kSignMask = pset1<Packet>(-Scalar(0));
+  const Packet kZero = pzero(x);
+  const Packet kOne = pset1<Packet>(Scalar(1));
+  const Packet kPi = pset1<Packet>(Scalar(EIGEN_PI));
+  const Packet kInf = pset1<Packet>(NumTraits<Scalar>::infinity());
+
+  const Packet abs_x = pabs(x);
+  const Packet x_is_zero = pcmp_eq(abs_x, kZero);
+  const Packet x_is_inf = pcmp_eq(abs_x, kInf);
+  const Packet x_has_signbit = psignbit(x);
+
+  const Packet abs_y = pabs(y);
+  const Packet y_is_zero = pcmp_eq(abs_y, kZero);
+  const Packet y_is_inf = pcmp_eq(abs_y, kInf);
+  const Packet y_signmask = pand(y, kSignMask);
+
+  const Packet arg_signmask = pand(pxor(x, y), kSignMask);
+  const Packet shift = pxor(pand(x_has_signbit, kPi), y_signmask);
+
+  // bend two rules:
+  // 1) 0 / 0 == 0
+  // 2) inf / inf == 1
+  // otherwise, evaluate atan(y/x) as usual and shift to the appropriate quadrant
+
+  Packet arg = pdiv(abs_y, abs_x);
+  arg = pselect(pand(x_is_zero, y_is_zero), kZero, arg);
+  arg = pselect(pand(x_is_inf, y_is_inf), kOne, arg);
+
+  Packet result = patan(arg);
+  result = pxor(result, arg_signmask);
+  result = padd(result, shift);
+
+  return result;
+}
+
 } // end namespace internal
 
 } // end namespace Eigen