Compute float64 exp, sin, cos and erf in double

mlx/backend/cpu/simd/math.h

@@ -8,6 +8,21 @@ namespace mlx::core::simd {
 
 constexpr float inf = std::numeric_limits<float>::infinity();
 
+/**
+ * Evaluate a scalar function on each lane.
+ *
+ * The approximations in this file are tuned for float32. They are not accurate
+ * enough for float64, so double precision falls back to libm through here.
+ */
+template <typename T, int N, typename Op>
+Simd<T, N> lanewise(Simd<T, N> in, Op op) {
+  Simd<T, N> out;
+  for (int i = 0; i < N; ++i) {
+    out[i] = op(in[i]);
+  }
+  return out;
+}
+
 /**
  * Compute exp(x) in an optimizer friendly way as follows:
  *
@@ -28,6 +43,10 @@ template <typename T, int N>
 Simd<T, N> exp(Simd<T, N> in) {
   if constexpr (is_complex<T>) {
     return Simd<T, 1>{std::exp(in.value)};
+  } else if constexpr (std::is_same_v<T, double>) {
+    // The polynomial below also saturates outside of float's range, which
+    // would clamp double results to zero or infinity beyond +-88.
+    return lanewise(in, [](double v) { return std::exp(v); });
   } else {
     Simd<float, N> x_init = in;
     auto x = x_init * 1.442695f; // multiply with log_2(e)
@@ -119,6 +138,8 @@ template <typename T, int N>
 Simd<T, N> sin(Simd<T, N> x) {
   if constexpr (is_complex<T>) {
     return std::sin(x.value);
+  } else if constexpr (std::is_same_v<T, double>) {
+    return lanewise(x, [](double v) { return std::sin(v); });
   } else {
     return sincos<true>(x);
   }
@@ -128,23 +149,29 @@ template <typename T, int N>
 Simd<T, N> cos(Simd<T, N> x) {
   if constexpr (is_complex<T>) {
     return std::cos(x.value);
+  } else if constexpr (std::is_same_v<T, double>) {
+    return lanewise(x, [](double v) { return std::cos(v); });
   } else {
     return sincos<false>(x);
   }
 }
 
 template <typename T, int N>
 Simd<T, N> erf(Simd<T, N> x) {
-  // https://github.com/pytorch/pytorch/blob/abf28982a8cb43342e7669d859de9543fd804cc9/aten/src/ATen/cpu/vec/vec256/vec256_float.h#L175
-  Simd<float, N> v = x;
-  auto t = recip(fma(Simd<float, N>(0.3275911f), abs(v), 1.0f));
-  auto r = fma(Simd<float, N>(1.061405429f), t, -1.453152027f);
-  r = fma(r, t, 1.421413741f);
-  r = fma(r, t, -0.284496736f);
-  r = fma(r, t, 0.254829592f);
-  auto e = -exp(-v * v);
-  auto result = Simd<T, N>(fma(e * t, r, 1.0f));
-  return select(x > 0, result, -result);
+  if constexpr (std::is_same_v<T, double>) {
+    return lanewise(x, [](double v) { return std::erf(v); });
+  } else {
+    // https://github.com/pytorch/pytorch/blob/abf28982a8cb43342e7669d859de9543fd804cc9/aten/src/ATen/cpu/vec/vec256/vec256_float.h#L175
+    Simd<float, N> v = x;
+    auto t = recip(fma(Simd<float, N>(0.3275911f), abs(v), 1.0f));
+    auto r = fma(Simd<float, N>(1.061405429f), t, -1.453152027f);
+    r = fma(r, t, 1.421413741f);
+    r = fma(r, t, -0.284496736f);
+    r = fma(r, t, 0.254829592f);
+    auto e = -exp(-v * v);
+    auto result = Simd<T, N>(fma(e * t, r, 1.0f));
+    return select(x > 0, result, -result);
+  }
 }
 
 template <typename T, int N>

python/tests/test_double.py

@@ -301,6 +301,42 @@ def test_linspace(self):
             vals = mx.linspace(0, math.pi, 2, mx.float64)
             self.assertEqual(vals.tolist()[1], math.pi)
 
+    def test_transcendental_precision(self):
+        # exp, sin, cos and erf use a fast float32 polynomial on the CPU. For
+        # float64 inputs they must be evaluated in double precision, otherwise
+        # the result silently carries only float32 accuracy.
+        with mx.stream(mx.cpu):
+            x_np = np.linspace(-3.0, 3.0, 257, dtype=np.float64)
+            x = mx.array(x_np, dtype=mx.float64)
+
+            erf_np = np.vectorize(math.erf, otypes=[np.float64])
+            for mx_op, np_op in [
+                (mx.exp, np.exp),
+                (mx.sin, np.sin),
+                (mx.cos, np.cos),
+                (mx.erf, erf_np),
+            ]:
+                y = np.array(mx_op(x), copy=False)
+                self.assertEqual(y.dtype, np.float64)
+                self.assertTrue(
+                    np.allclose(y, np_op(x_np), rtol=1e-14, atol=1e-300),
+                    msg=f"{mx_op.__name__} on float64 is not double precision",
+                )
+
+    def test_exp_range(self):
+        # The float32 implementation saturates outside [-88, 88]. A float64
+        # exp must cover the whole double range instead.
+        with mx.stream(mx.cpu):
+            x_np = np.array([-700.0, -200.0, -89.0, 89.0, 200.0, 700.0])
+            y = np.array(mx.exp(mx.array(x_np, dtype=mx.float64)), copy=False)
+            self.assertTrue(np.all(np.isfinite(y)))
+            self.assertTrue(np.all(y > 0))
+            self.assertTrue(np.allclose(y, np.exp(x_np), rtol=1e-14))
+
+            # float32 keeps saturating, as before.
+            y32 = np.array(mx.exp(mx.array(x_np, dtype=mx.float32)), copy=False)
+            self.assertTrue(np.isinf(y32[-1]))
+
 
 if __name__ == "__main__":
     mlx_tests.MLXTestRunner()