CPU-Strided-Complex Fixes for real and imag ops (pytorch#29840)

Summary:
In-tree changes to pytorch to support complex numbers are being submitted here.
Out-of-tree support for complex numbers is here: [pytorch-cpu-strided-complex extension](https://gitlab.com/pytorch-complex/pytorch-cpu-strided-complex)

- [x]  Replaced std:real(a) with a.real() in kernel level code.
- [x]  Fixed Vec256_base implementation of complex ops so that it works correctly on Non-AVX devices.
- [x]  Fix NumericUtils.h

cc: iotamudelta, ezyang, bddppq, zasdfgbnm
Pull Request resolved: pytorch#29840

Differential Revision: D18531274

Pulled By: ezyang

fbshipit-source-id: 0fa842c68e4bd55134fe0271880e2d15fe692b7f

Author: dylanbespalko

aten/src/ATen/NumericUtils.h

@@ -34,9 +34,9 @@ inline C10_HOST_DEVICE bool _isnan(T val) {
 }
 
 template <typename T,
-          typename std::enable_if<std::is_complex_t<T>::value, int>::type = 0>
+          typename std::enable_if<c10::is_complex_t<T>::value, int>::type = 0>
 inline bool _isnan(T val) {
-  return std::isnan(std::real(val)) || std::isnan(std::imag(val));
+  return std::isnan(val.real()) || std::isnan(val.imag());
 }
 
 inline C10_HOST_DEVICE bool _isnan(at::BFloat16 val) {

aten/src/ATen/cpu/vec256/vec256_base.h

@@ -179,42 +179,86 @@ struct Vec256 {
     }
     return ret;
   }
-  template <typename other_t = T,
-            typename std::enable_if<!std::is_floating_point<other_t>::value && !std::is_complex_t<other_t>::value, int>::type = 0>
+  template <typename other_t_abs = T,
+            typename std::enable_if<!std::is_floating_point<other_t_abs>::value && !c10::is_complex_t<other_t_abs>::value, int>::type = 0>
   Vec256<T> abs() const {
-    // other_t is for SFINAE and clarity. Make sure it is not changed.
-    static_assert(std::is_same<other_t, T>::value, "other_t must be T");
-    return map([](T x) -> T { return x < static_cast<other_t>(0) ? -x : x; });
+    // other_t_abs is for SFINAE and clarity. Make sure it is not changed.
+    static_assert(std::is_same<other_t_abs, T>::value, "other_t_abs must be T");
+    return map([](T x) -> T { return x < static_cast<T>(0) ? -x : x; });
   }
-  template <typename float_t = T,
-            typename std::enable_if<std::is_floating_point<float_t>::value, int>::type = 0>
+  template <typename float_t_abs = T,
+            typename std::enable_if<std::is_floating_point<float_t_abs>::value, int>::type = 0>
   Vec256<T> abs() const {
-    // float_t is for SFINAE and clarity. Make sure it is not changed.
-    static_assert(std::is_same<float_t, T>::value, "float_t must be T");
+    // float_t_abs is for SFINAE and clarity. Make sure it is not changed.
+    static_assert(std::is_same<float_t_abs, T>::value, "float_t_abs must be T");
     // Specifically deal with floating-point because the generic code above won't handle -0.0 (which should result in
     // 0.0) properly.
     return map(std::abs);
   }
-  template <typename complex_t = T,
-            typename std::enable_if<std::is_complex_t<complex_t>::value, int>::type = 0>
+  template <typename complex_t_abs = T,
+            typename std::enable_if<c10::is_complex_t<complex_t_abs>::value, int>::type = 0>
   Vec256<T> abs() const {
-    // complex_t is for SFINAE and clarity. Make sure it is not changed.
-    static_assert(std::is_same<complex_t, T>::value, "complex_t must be T");
+    // complex_t_abs is for SFINAE and clarity. Make sure it is not changed.
+    static_assert(std::is_same<complex_t_abs, T>::value, "complex_t_abs must be T");
     // Specifically map() does not perform the type conversion needed by abs.
-    return map([](T x) { return (T)std::abs(x); });
+    return map([](T x) { return static_cast<T>(std::abs(x)); });
   }
+  template <typename other_t_angle = T,
+            typename std::enable_if<!c10::is_complex_t<other_t_angle>::value, int>::type = 0>
   Vec256<T> angle() const {
-    return *this;
+    // other_t_angle is for SFINAE and clarity. Make sure it is not changed.
+    static_assert(std::is_same<other_t_angle, T>::value, "other_t_angle must be T");
+    return Vec256(0);
+  }
+  template <typename complex_t_angle = T,
+            typename std::enable_if<c10::is_complex_t<complex_t_angle>::value, int>::type = 0>
+  Vec256<T> angle() const {
+    // complex_t_angle is for SFINAE and clarity. Make sure it is not changed.
+    static_assert(std::is_same<complex_t_angle, T>::value, "complex_t_angle must be T");
+    return map([](T x) { return static_cast<T>(std::arg(x)); });
   }
+  template <typename other_t_real = T,
+            typename std::enable_if<!c10::is_complex_t<other_t_real>::value, int>::type = 0>
   Vec256<T> real() const {
+    // other_t_real is for SFINAE and clarity. Make sure it is not changed.
+    static_assert(std::is_same<other_t_real, T>::value, "other_t_real must be T");
     return *this;
   }
+  template <typename complex_t_real = T,
+            typename std::enable_if<c10::is_complex_t<complex_t_real>::value, int>::type = 0>
+  Vec256<T> real() const {
+    // complex_t_real is for SFINAE and clarity. Make sure it is not changed.
+    static_assert(std::is_same<complex_t_real, T>::value, "complex_t_real must be T");
+    return map([](T x) { return static_cast<T>(x.real()); });
+  }
+  template <typename other_t_imag = T,
+            typename std::enable_if<!c10::is_complex_t<other_t_imag>::value, int>::type = 0>
   Vec256<T> imag() const {
-    return *this;
+    // other_t_imag is for SFINAE and clarity. Make sure it is not changed.
+    static_assert(std::is_same<other_t_imag, T>::value, "other_t_imag must be T");
+    return Vec256(0);
+  }
+  template <typename complex_t_imag = T,
+            typename std::enable_if<c10::is_complex_t<complex_t_imag>::value, int>::type = 0>
+  Vec256<T> imag() const {
+    // complex_t_imag is for SFINAE and clarity. Make sure it is not changed.
+    static_assert(std::is_same<complex_t_imag, T>::value, "complex_t_imag must be T");
+    return map([](T x) { return static_cast<T>(x.imag()); });
   }
+  template <typename other_t_conj = T,
+            typename std::enable_if<!c10::is_complex_t<other_t_conj>::value, int>::type = 0>
   Vec256<T> conj() const {
+    // other_t_conj is for SFINAE and clarity. Make sure it is not changed.
+    static_assert(std::is_same<other_t_conj, T>::value, "other_t_conj must be T");
     return *this;
   }
+  template <typename complex_t_conj = T,
+            typename std::enable_if<c10::is_complex_t<complex_t_conj>::value, int>::type = 0>
+  Vec256<T> conj() const {
+    // complex_t_conj is for SFINAE and clarity. Make sure it is not changed.
+    static_assert(std::is_same<complex_t_conj, T>::value, "complex_t_conj must be T");
+    return map([](T x) { return static_cast<T>(std::conj(x)); });
+  }
   Vec256<T> acos() const {
     return map(std::acos);
   }
@@ -259,14 +303,14 @@ struct Vec256 {
     return map(std::log1p);
   }
   template <typename other_t_log2 = T,
-            typename std::enable_if<!std::is_complex_t<other_t_log2>::value, int>::type = 0>
+            typename std::enable_if<!c10::is_complex_t<other_t_log2>::value, int>::type = 0>
   Vec256<T> log2() const {
     // other_t_log2 is for SFINAE and clarity. Make sure it is not changed.
     static_assert(std::is_same<other_t_log2, T>::value, "other_t_log2 must be T");
     return map(std::log2);
   }
   template <typename complex_t_log2 = T,
-            typename std::enable_if<std::is_complex_t<complex_t_log2>::value, int>::type = 0>
+            typename std::enable_if<c10::is_complex_t<complex_t_log2>::value, int>::type = 0>
   Vec256<T> log2() const {
     // complex_t_log2 is for SFINAE and clarity. Make sure it is not changed.
     static_assert(std::is_same<complex_t_log2, T>::value, "complex_t_log2 must be T");
@@ -395,7 +439,7 @@ template <class T> Vec256<T> inline operator||(
 // Implements the IEEE 754 201X `maximum` operation, which propagates NaN if
 // either input is a NaN.
 template <class T,
-          typename std::enable_if<!std::is_complex_t<T>::value, int>::type = 0>
+          typename std::enable_if<!c10::is_complex_t<T>::value, int>::type = 0>
 Vec256<T> inline maximum(const Vec256<T> &a, const Vec256<T> &b) {
   Vec256<T> c = Vec256<T>();
   for (int i = 0; i != Vec256<T>::size(); i++) {
@@ -411,7 +455,7 @@ Vec256<T> inline maximum(const Vec256<T> &a, const Vec256<T> &b) {
 }
 
 template <class T,
-          typename std::enable_if<std::is_complex_t<T>::value, int>::type = 0>
+          typename std::enable_if<c10::is_complex_t<T>::value, int>::type = 0>
 Vec256<T> inline maximum(const Vec256<T> &a, const Vec256<T> &b) {
   Vec256<T> c = Vec256<T>();
   for (int i = 0; i != Vec256<T>::size(); i++) {
@@ -438,7 +482,7 @@ inline T maximum(const T& a, const T& b) {
 // Implements the IEEE 754 201X `minimum` operation, which propagates NaN if
 // either input is a NaN.
 template <class T,
-          typename std::enable_if<!std::is_complex_t<T>::value, int>::type = 0>
+          typename std::enable_if<!c10::is_complex_t<T>::value, int>::type = 0>
 Vec256<T> inline minimum(const Vec256<T> &a, const Vec256<T> &b) {
   Vec256<T> c = Vec256<T>();
   for (int i = 0; i != Vec256<T>::size(); i++) {
@@ -454,7 +498,7 @@ Vec256<T> inline minimum(const Vec256<T> &a, const Vec256<T> &b) {
 }
 
 template <class T,
-          typename std::enable_if<std::is_complex_t<T>::value, int>::type = 0>
+          typename std::enable_if<c10::is_complex_t<T>::value, int>::type = 0>
 Vec256<T> inline minimum(const Vec256<T> &a, const Vec256<T> &b) {
   Vec256<T> c = Vec256<T>();
   for (int i = 0; i != Vec256<T>::size(); i++) {
@@ -480,7 +524,7 @@ inline T minimum(const T& a, const T& b) {
 
 // To save BC, it will not propagate NaN based on IEEE 754 201X
 template <class T,
-          typename std::enable_if<!std::is_complex_t<T>::value, int>::type = 0>
+          typename std::enable_if<!c10::is_complex_t<T>::value, int>::type = 0>
 Vec256<T> inline clamp(const Vec256<T> &a, const Vec256<T> &min_vec, const Vec256<T> &max_vec) {
   Vec256<T> c = Vec256<T>();
   for (int i = 0; i != Vec256<T>::size(); i++) {
@@ -490,7 +534,7 @@ Vec256<T> inline clamp(const Vec256<T> &a, const Vec256<T> &min_vec, const Vec25
 }
 
 template <class T,
-          typename std::enable_if<std::is_complex_t<T>::value, int>::type = 0>
+          typename std::enable_if<c10::is_complex_t<T>::value, int>::type = 0>
 Vec256<T> inline clamp(const Vec256<T> &a, const Vec256<T> &min_vec, const Vec256<T> &max_vec) {
   Vec256<T> c = Vec256<T>();
   for (int i = 0; i != Vec256<T>::size(); i++) {
@@ -500,7 +544,7 @@ Vec256<T> inline clamp(const Vec256<T> &a, const Vec256<T> &min_vec, const Vec25
 }
 
 template <class T,
-          typename std::enable_if<!std::is_complex_t<T>::value, int>::type = 0>
+          typename std::enable_if<!c10::is_complex_t<T>::value, int>::type = 0>
 Vec256<T> inline clamp_max(const Vec256<T> &a, const Vec256<T> &max_vec) {
   Vec256<T> c = Vec256<T>();
   for (int i = 0; i != Vec256<T>::size(); i++) {
@@ -510,7 +554,7 @@ Vec256<T> inline clamp_max(const Vec256<T> &a, const Vec256<T> &max_vec) {
 }
 
 template <class T,
-          typename std::enable_if<std::is_complex_t<T>::value, int>::type = 0>
+          typename std::enable_if<c10::is_complex_t<T>::value, int>::type = 0>
 Vec256<T> inline clamp_max(const Vec256<T> &a, const Vec256<T> &max_vec) {
   Vec256<T> c = Vec256<T>();
   for (int i = 0; i != Vec256<T>::size(); i++) {
@@ -520,7 +564,7 @@ Vec256<T> inline clamp_max(const Vec256<T> &a, const Vec256<T> &max_vec) {
 }
 
 template <class T,
-          typename std::enable_if<!std::is_complex_t<T>::value, int>::type = 0>
+          typename std::enable_if<!c10::is_complex_t<T>::value, int>::type = 0>
 Vec256<T> inline clamp_min(const Vec256<T> &a, const Vec256<T> &min_vec) {
   Vec256<T> c = Vec256<T>();
   for (int i = 0; i != Vec256<T>::size(); i++) {
@@ -530,7 +574,7 @@ Vec256<T> inline clamp_min(const Vec256<T> &a, const Vec256<T> &min_vec) {
 }
 
 template <class T,
-          typename std::enable_if<std::is_complex_t<T>::value, int>::type = 0>
+          typename std::enable_if<c10::is_complex_t<T>::value, int>::type = 0>
 Vec256<T> inline clamp_min(const Vec256<T> &a, const Vec256<T> &min_vec) {
   Vec256<T> c = Vec256<T>();
   for (int i = 0; i != Vec256<T>::size(); i++) {

aten/src/ATen/cpu/vec256/vec256_complex_double.h

@@ -24,14 +24,14 @@ template <> class Vec256<std::complex<double>> {
   Vec256() {}
   Vec256(__m256d v) : values(v) {}
   Vec256(std::complex<double> val) {
-    double real_value = std::real(val);
-    double imag_value = std::imag(val);
+    double real_value = val.real();
+    double imag_value = val.imag();
     values = _mm256_setr_pd(real_value, imag_value,
                             real_value, imag_value);
   }
   Vec256(std::complex<double> val1, std::complex<double> val2) {
-    values = _mm256_setr_pd(std::real(val1), std::imag(val1),
-                            std::real(val2), std::imag(val2));
+    values = _mm256_setr_pd(val1.real(), val1.imag(),
+                            val2.real(), val2.imag());
   }
   operator __m256d() const {
     return values;

aten/src/ATen/cpu/vec256/vec256_complex_float.h

@@ -24,19 +24,19 @@ template <> class Vec256<std::complex<float>> {
   Vec256() {}
   Vec256(__m256 v) : values(v) {}
   Vec256(std::complex<float> val) {
-    float real_value = std::real(val);
-    float imag_value = std::imag(val);
+    float real_value = val.real();
+    float imag_value = val.imag();
     values = _mm256_setr_ps(real_value, imag_value,
                             real_value, imag_value,
                             real_value, imag_value,
                             real_value, imag_value
                             );
   }
   Vec256(std::complex<float> val1, std::complex<float> val2, std::complex<float> val3, std::complex<float> val4) {
-    values = _mm256_setr_ps(std::real(val1), std::imag(val1),
-                            std::real(val2), std::imag(val2),
-                            std::real(val3), std::imag(val3),
-                            std::real(val4), std::imag(val4)
+    values = _mm256_setr_ps(val1.real(), val1.imag(),
+                            val2.real(), val2.imag(),
+                            val3.real(), val3.imag(),
+                            val4.real(), val4.imag()
                             );
   }
   operator __m256() const {