Compile float64 functions on CPU (#2311)

mlx/backend/common/compiled.cpp

@@ -14,6 +14,8 @@ void print_constant(std::ostream& os, const array& x) {
       return print_float_constant<float16_t>(os, x);
     case bfloat16:
       return print_float_constant<bfloat16_t>(os, x);
+    case float64:
+      return print_float_constant<double>(os, x);
     case complex64:
       return print_complex_constant<complex64_t>(os, x);
     case int8:
@@ -50,6 +52,8 @@ std::string get_type_string(Dtype d) {
       return "float16_t";
     case bfloat16:
       return "bfloat16_t";
+    case float64:
+      return "double";
     case complex64:
       return "complex64_t";
     case bool_:

mlx/backend/common/compiled.h

@@ -18,8 +18,12 @@ std::string get_type_string(Dtype d);
 template <typename T>
 void print_float_constant(std::ostream& os, const array& x) {
   auto old_precision = os.precision();
-  os << std::setprecision(std::numeric_limits<float>::digits10 + 1)
+  if constexpr (std::is_same_v<T, double>) {
-     << x.item<T>() << std::setprecision(old_precision);
+    os << std::setprecision(std::numeric_limits<double>::digits10 + 1);
+  } else {
+    os << std::setprecision(std::numeric_limits<float>::digits10 + 1);
+  }
+  os << x.item<T>() << std::setprecision(old_precision);
 }
 
 template <typename T>

python/src/convert.cpp

@@ -205,6 +205,8 @@ nb::object to_scalar(mx::array& a) {
       return nb::cast(static_cast<float>(a.item<mx::bfloat16_t>()));
     case mx::complex64:
       return nb::cast(a.item<std::complex<float>>());
+    case mx::float64:
+      return nb::cast(a.item<double>());
     default:
       throw nb::type_error("type cannot be converted to Python scalar.");
   }

python/tests/test_compile.py

@@ -2,6 +2,7 @@
 
 import gc
 import io
+import math
 import unittest
 from functools import partial
 
@@ -979,6 +980,17 @@ class TestCompile(mlx_tests.MLXTestCase):
 
         self.assertEqual(mem_pre, mem_post)
 
+    def test_double_constant(self):
+        with mx.stream(mx.cpu):
+            x = mx.array(1.0, dtype=mx.float64)
+
+            def fun(x):
+                return (x + math.pi) * 2.0
+
+            y = fun(x).item()
+            y_compiled = mx.compile(fun)(x).item()
+            self.assertEqual(y, y_compiled)
+
 
 if __name__ == "__main__":
     mlx_tests.MLXTestRunner()