Use c10 math constants consistently in Math.h (pytorch#91967)

On MSVC the `M_` constants are hidden behind the `USE_MATH_DEFINES` macro, so it's better to avoid them in headers otherwise the include order can break compilation. Pull Request resolved: pytorch#91967 Approved by: https://github.com/malfet
Mookel · Jan 12, 2023 · 8acf0e6 · 8acf0e6
1 parent c7a22bb
commit 8acf0e6
Show file tree

Hide file tree

Showing 3 changed files with 13 additions and 2 deletions.
diff --git a/aten/src/ATen/native/Math.h b/aten/src/ATen/native/Math.h
@@ -2169,7 +2169,7 @@ calc_erfcx(T x)
  */
 template <typename T>
 static inline C10_HOST_DEVICE T calc_log_ndtr(T x) {
-  T t = x * M_SQRT1_2;
+  T t = x * c10::frac_sqrt_2<T>;
   if (x < T{-1.0}) {
     return std::log(calc_erfcx(-t) / 2) - t * t;
   } else {
@@ -2295,7 +2295,7 @@ static inline C10_HOST_DEVICE T airy_ai_forward(T x) {
             agd = agd * (z * z) + AGD[index];
         }
 
-        T t = T(-2.0) * x * std::sqrt(-x) / T(3.0) + T(0.25) * M_PI;
+        T t = T(-2.0) * x * std::sqrt(-x) / T(3.0) + T(0.25) * c10::pi<T>;
 
         return T(5.64189583547756286948e-01) / std::sqrt(std::sqrt(-x)) * (std::sin(t) * (T(1.0) + z * z * afn / afd) - std::cos(t) * (z * agn / agd));
     }

diff --git a/c10/util/MathConstants.cpp b/c10/util/MathConstants.cpp
@@ -11,3 +11,6 @@
 #include <math.h>
 
 static_assert(M_PI == c10::pi<double>, "c10::pi<double> must be equal to M_PI");
+static_assert(
+    M_SQRT1_2 == c10::frac_sqrt_2<double>,
+    "c10::frac_sqrt_2<double> must be equal to M_SQRT1_2");
diff --git a/c10/util/MathConstants.h b/c10/util/MathConstants.h
@@ -32,6 +32,11 @@ C10_HOST_DEVICE inline constexpr T frac_1_sqrt_pi() {
   return static_cast<T>(0.564189583547756286948079451560772);
 }
 
+template <typename T>
+C10_HOST_DEVICE inline constexpr T frac_sqrt_2() {
+  return static_cast<T>(0.707106781186547524400844362104849);
+}
+
 template <typename T>
 C10_HOST_DEVICE inline constexpr T frac_sqrt_3() {
   return static_cast<T>(0.577350269189625764509148780501957);
@@ -103,6 +108,9 @@ constexpr T frac_1_pi = c10::detail::frac_1_pi<T>();
 template <typename T>
 constexpr T frac_1_sqrt_pi = c10::detail::frac_1_sqrt_pi<T>();
 
+template <typename T>
+constexpr T frac_sqrt_2 = c10::detail::frac_sqrt_2<T>();
+
 template <typename T>
 constexpr T frac_sqrt_3 = c10::detail::frac_sqrt_3<T>();