vsx: complex types use std::sqrt

quickwritereader · quickwritereader · commit 97c8e17a521b · 2020-12-09T00:01:57.000Z
diff --git a/aten/src/ATen/cpu/vec256/vsx/vec256_complex_double_vsx.h b/aten/src/ATen/cpu/vec256/vsx/vec256_complex_double_vsx.h
@@ -364,17 +364,7 @@ class Vec256<ComplexDbl> {
   }
 
   Vec256<ComplexDbl> sqrt() const {
-    //   sqrt(a + bi)
-    // = sqrt(2)/2 * [sqrt(sqrt(a**2 + b**2) + a) + sgn(b)*sqrt(sqrt(a**2 +
-    // b**2) - a)i] = sqrt(2)/2 * [sqrt(abs() + a) + sgn(b)*sqrt(abs() - a)i]
-
-    auto sign = *this & vd_isign_mask;
-    auto factor = sign | vd_sqrt2_2;
-    auto a_a = el_mergee();
-    // a_a.dump();
-    a_a = a_a ^ vd_isign_mask; // a -a
-    auto res_re_im = (abs_() + a_a).elwise_sqrt(); // sqrt(abs + a) sqrt(abs - a)
-    return factor.elwise_mult(res_re_im);
+    return map(std::sqrt);  
   }
 
   Vec256<ComplexDbl> reciprocal() const {
diff --git a/aten/src/ATen/cpu/vec256/vsx/vec256_complex_float_vsx.h b/aten/src/ATen/cpu/vec256/vsx/vec256_complex_float_vsx.h
@@ -417,18 +417,7 @@ class Vec256<ComplexFlt> {
   }
 
   Vec256<ComplexFlt> sqrt() const {
-    //   sqrt(a + bi)
-    // = sqrt(2)/2 * [sqrt(sqrt(a**2 + b**2) + a) + sgn(b)*sqrt(sqrt(a**2 +
-    // b**2) - a)i] = sqrt(2)/2 * [sqrt(abs() + a) + sgn(b)*sqrt(abs() - a)i]
-
-    auto sign = *this & isign_mask;
-    auto factor = sign | sqrt2_2;
-    auto a_a = el_mergee();
-    // a_a.dump();
-    a_a = a_a ^ isign_mask; // a -a
-    auto res_re_im =
-        (abs_() + a_a).elwise_sqrt(); // sqrt(abs + a) sqrt(abs - a)
-    return factor.elwise_mult(res_re_im);
+    return map(std::sqrt);  
   }
 
   Vec256<ComplexFlt> reciprocal() const {
