Repository: tinygrad/tinygrad
Pull request: F64 support
Adds float64 support to the different runtimes, also included an excluded field on dtypes for unsupported platforms

From my testing float64 is not supported on metal and OpenCL on M2 Mac. OpenCL float64 is fine on my windows machine although this is an optional feature and ideally should check the OpenCL extensions for it.

Also brings the onnx stat line to:
190 failed, 650 passed, 1794 skipped, 2 warnings in 14.73s

Tested on:
- CPU
- GPU
- CUDA
- LLVM
- CLANG

Closes #891 

Problem statement:
float64 support
Support for float64 over all plattforms.
This would also fix up to 20 ONNX tests.

Is there a plan to add this? Otherwise there is no way to fully support ONNX, since some tests are hard relying on that.

Hints from issue discussion:
Metal currently does not support float64 / double
Would merge PR to skip those tests. float64 is not in scope for tinygrad.
You did not :-D
There are just 20 random tests in the ONNX test suite that require float64. To me its surprising, that we only have 20, since all operations support float64s.

I fear that tinygrad needs a better engineered version of the ONNX test suite to cover the parts tinygrad supports and ignores the others.

In any case we should close this issue, since it cannot be resolved.

That PR didn't make it clear that those tests needed float64. Idk, if they are really needed, maybe we should support it, at least on some backends.

Commit message: Merge branch 'f64' of github.com:dc-dc-dc/tinygrad into f64

Patch to review:
diff --git a/test/test_dtype.py b/test/test_dtype.py
index 65fde81ab999b..8fb1202c5252a 100644
--- a/test/test_dtype.py
+++ b/test/test_dtype.py
@@ -52,6 +52,17 @@ def test_half_mul_upcast_float(self): _test_mul_upcast(Tensor([1,2,3,4], dtype=d
   def test_half_matmul_upcast_float(self): _test_matmul_upcast(Tensor([[1,2],[3,4]], dtype=dtypes.float16), Tensor.eye(2, dtype=dtypes.float32), dtypes.float32, [[1,2],[3,4]])
   def test_int8_matmul_upcast_half(self): _test_matmul_upcast(Tensor([[1,2],[3,4]], dtype=dtypes.int8), Tensor.eye(2, dtype=dtypes.float16), dtypes.float16, [[1,2],[3,4]])
 
+class TestFloatDtype(unittest.TestCase):
+  def test_float64_to_np(self): _test_to_np(Tensor([1,2,3,4], dtype=dtypes.float64), np.float64, [1,2,3,4])
+  def test_float64_add(self): _test_add(Tensor([1,2,3,4], dtype=dtypes.float64), Tensor([1,2,3,4], dtype=dtypes.float64), dtypes.float64, [2,4,6,8])
+  def test_float64_mul(self): _test_mul(Tensor([1,2,3,4], dtype=dtypes.float64), Tensor([1,2,3,4], dtype=dtypes.float64), dtypes.float64, [1,4,9,16])
+  def test_float64_matmul(self): _test_matmul(Tensor([[1,2],[3,4]], dtype=dtypes.float64), Tensor.eye(2, dtype=dtypes.float64), dtypes.float64, [[1,2],[3,4]])
+  def test_float64_cast_float(self): _test_cast(Tensor([1,2]), dtypes.float64, [1,2])
+
+  def test_float_add_upcast_float64(self): _test_add_upcast(Tensor([1,2,3,4], dtype=dtypes.float32), Tensor([1,2,3,4], dtype=dtypes.float64), dtypes.float64, [2,4,6,8])
+  def test_float_mul_upcast_float64(self): _test_mul_upcast(Tensor([1,2,3,4], dtype=dtypes.float32), Tensor([1,2,3,4], dtype=dtypes.float64), dtypes.float64, [1,4,9,16])
+  def test_float_matmul_upcast_float64(self): _test_matmul_upcast(Tensor([[1,2],[3,4]], dtype=dtypes.float32), Tensor.eye(2, dtype=dtypes.float64), dtypes.float64, [[1,2],[3,4]])
+
 class TestInt8Dtype(unittest.TestCase):
   def test_int8_to_np(self): _test_to_np(Tensor([1,2,3,4], dtype=dtypes.int8), np.int8, [1,2,3,4])
   def test_uint8_to_np(self): _test_to_np(Tensor([1,2,3,4], dtype=dtypes.uint8), np.uint8, [1,2,3,4])
diff --git a/tinygrad/codegen/llvmir.py b/tinygrad/codegen/llvmir.py
index 4fc21b545ed82..7f0c5d315a8a3 100644
--- a/tinygrad/codegen/llvmir.py
+++ b/tinygrad/codegen/llvmir.py
@@ -37,7 +37,7 @@ def uops_to_llvm_ir(uops:List[UOp], bufs:List[LazyBuffer]) -> str:
   module = ir.Module(name=__file__)
 
   # create llvm function
-  func_dtypes = [{dtypes.float16:ir.HalfType(), dtypes.float32:ir.FloatType(), dtypes.int8:ir.IntType(8), dtypes.uint8:ir.IntType(8), dtypes.bool: ir.IntType(1), dtypes.int64: ir.IntType(64)}[buf.dtype] for buf in bufs]
+  func_dtypes = [{dtypes.float16:ir.HalfType(), dtypes.float32:ir.FloatType(), dtypes.float64:ir.DoubleType(), dtypes.int8:ir.IntType(8), dtypes.uint8:ir.IntType(8), dtypes.bool: ir.IntType(1), dtypes.int64: ir.IntType(64)}[buf.dtype] for buf in bufs]
   func = ir.Function(module, ir.FunctionType(ir.VoidType(), [x.as_pointer() for x in func_dtypes]), name='exec')
 
   # force llvmlite to allow us to add function attribute then add the attribute
@@ -90,6 +90,8 @@ def uops_to_llvm_ir(uops:List[UOp], bufs:List[LazyBuffer]) -> str:
       if func_dtypes[args.i] != ir.FloatType():
         if dtypes.is_int(bufs[args.i].dtype):
           val = bb[-1].uitofp(val, ir.FloatType()) if dtypes.is_unsigned(bufs[args.i].dtype) else bb[-1].sitofp(val, ir.FloatType())
+        elif bufs[args.i].dtype == dtypes.float64:
+          val = bb[-1].fptrunc(val, ir.FloatType())
         else:
           val = bb[-1].fpext(val, ir.FloatType())
       lvars[newvar] = val
@@ -100,6 +102,8 @@ def uops_to_llvm_ir(uops:List[UOp], bufs:List[LazyBuffer]) -> str:
       if func_dtypes[0] != ir.FloatType():
         if dtypes.is_int(bufs[args.i].dtype):
           element = bb[-1].fptoui(element, func_dtypes[0]) if dtypes.is_unsigned(bufs[args.i].dtype) else bb[-1].fptosi(element, func_dtypes[0])
+        elif bufs[args.i].dtype == dtypes.float64:
+          element = bb[-1].fpext(element, func_dtypes[0])
         else:
           element = bb[-1].fptrunc(element, func_dtypes[0])
       bb[-1].store(element, bb[-1].gep(func.args[args.i], [idx], inbounds=True))
diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py
index 6201b82d4d0a1..5a0fcbb51fbc6 100644
--- a/tinygrad/helpers.py
+++ b/tinygrad/helpers.py
@@ -1,10 +1,12 @@
 from __future__ import annotations
+import platform
 from dataclasses import dataclass, asdict
 import os, math, functools, time
 import numpy as np
 from typing import Tuple, Union, List, NamedTuple, Final, Iterator, ClassVar, Optional, Callable, Any
 ShapeType = Tuple[int, ...]
 # NOTE: helpers is not allowed to import from anything else in tinygrad
+OSX = platform.system() == "Darwin"
 
 def dedup(x): return list(dict.fromkeys(x))   # retains list order
 def prod(x:Union[List[int], Tuple[int, ...]]) -> int: return math.prod(x)
@@ -53,6 +55,7 @@ class DType(NamedTuple):
   itemsize: int
   name: str
   np: type  # TODO: someday this will be removed with the "remove numpy" project
+  excluded: Tuple[str,...] = ()
   def __repr__(self): return f"dtypes.{self.name}"
 
 # dependent typing?
@@ -69,7 +72,7 @@ class dtypes:
   @staticmethod # static methds on top, or bool in the type info will refer to dtypes.bool
   def is_int(x: DType)-> bool: return x in (dtypes.int8, dtypes.uint8, dtypes.int32, dtypes.int64)
   @staticmethod
-  def is_float(x: DType) -> bool: return x in (dtypes.float16, dtypes.float32)
+  def is_float(x: DType) -> bool: return x in (dtypes.float16, dtypes.float32, dtypes.float64)
   @staticmethod
   def is_unsigned(x: DType) -> bool: return x in (dtypes.uint8)
   @staticmethod
@@ -77,6 +80,7 @@ def from_np(x) -> DType: return asdict(dtypes())[np.dtype(x).name]
   bool: Final[DType] = DType(0, 1, "bool", bool)
   float16: Final[DType] = DType(0, 2, "half", np.float16)
   float32: Final[DType] = DType(4, 4, "float", np.float32)
+  float64: Final[DType] = DType(5, 8, "double", np.float64, ("METAL", "GPU" if OSX else ""))
   int8: Final[DType] = DType(0, 1, "char", np.int8)
   int32: Final[DType] = DType(1, 4, "int", np.int32)
   int64: Final[DType] = DType(2, 8, "int64", np.int64)
diff --git a/tinygrad/runtime/lib.py b/tinygrad/runtime/lib.py
index 32aa00a380986..e986c5f338a01 100644
--- a/tinygrad/runtime/lib.py
+++ b/tinygrad/runtime/lib.py
@@ -37,7 +37,7 @@ def _copyin(self, x:np.ndarray) -> None: np.copyto(self.toCPU(), x.reshape(-1))
 
 # this one is simple enough that i moved it out of the runtimes
 class RawMallocBuffer(RawBufferMapped):
-  def __init__(self, size, dtype: DType): super().__init__(size, dtype, ({dtypes.float32: ctypes.c_float, dtypes.float16: ctypes.c_int16, dtypes.int8: ctypes.c_int8, dtypes.uint8: ctypes.c_uint8, dtypes.bool: ctypes.c_uint8, dtypes.int32: ctypes.c_int32, dtypes.int64: ctypes.c_int64}[dtype] * size)())
+  def __init__(self, size, dtype: DType): super().__init__(size, dtype, ({dtypes.float32: ctypes.c_float, dtypes.float16: ctypes.c_int16, dtypes.float64: ctypes.c_double, dtypes.int8: ctypes.c_int8, dtypes.uint8: ctypes.c_uint8, dtypes.bool: ctypes.c_uint8, dtypes.int32: ctypes.c_int32, dtypes.int64: ctypes.c_int64}[dtype] * size)())
   def _buffer(self): return memoryview(self._buf)
 
 class RawBufferCopyInOut(RawBufferCopyIn):
diff --git a/tinygrad/runtime/ops_gpu.py b/tinygrad/runtime/ops_gpu.py
index 2a027ad7bbc07..b347b0d4f483d 100644
--- a/tinygrad/runtime/ops_gpu.py
+++ b/tinygrad/runtime/ops_gpu.py
@@ -1,14 +1,13 @@
 from __future__ import annotations
-import platform, pathlib
+import pathlib
 import numpy as np
 import pyopencl as cl  # type: ignore
 from typing import Optional, List
-from tinygrad.helpers import DEBUG, getenv, prod, ImageDType
+from tinygrad.helpers import DEBUG, getenv, prod, ImageDType, OSX
 from tinygrad.ops import Compiled
 from tinygrad.runtime.lib import RawBufferCopyInOut
 from tinygrad.codegen.cstyle import CStyleCodegen, CStyleLanguage
 
-OSX = platform.system() == "Darwin"
 OSX_TIMING_RATIO = (125/3) if OSX else 1.0   # see test/external_osx_profiling.py to determine this ratio. it's in like GPU clocks or something
 FLOAT16 = getenv("FLOAT16", 0)
 
@@ -87,7 +86,7 @@ def __call__(self, global_size, local_size, *bufs, wait=False) -> Optional[float
 
 class CLCodegen(CStyleCodegen):
   lang = CStyleLanguage(
-    kernel_prefix = "#define int64 long\n__kernel", buffer_prefix = "__global ", smem_prefix = "__local ",
+    kernel_prefix = "#ifdef cl_khr_fp64\n#pragma OPENCL EXTENSION cl_khr_fp64 : enable\n#elif defined(cl_amd_fp64)\n#pragma OPENCL EXTENSION cl_amd_fp64 : enable\n#endif\n#define int64 long\n__kernel", buffer_prefix = "__global ", smem_prefix = "__local ",
     half_prekernel = "#pragma OPENCL EXTENSION cl_khr_fp16 : enable",
     barrier = "barrier(CLK_LOCAL_MEM_FENCE);", float4 = "(float4)",
     gid = [f'get_global_id({i})' for i in range(3)], lid = [f'get_local_id({i})' for i in range(3)], uses_vload=True)
diff --git a/tinygrad/runtime/ops_torch.py b/tinygrad/runtime/ops_torch.py
index 3320fc089f389..317443016ddda 100644
--- a/tinygrad/runtime/ops_torch.py
+++ b/tinygrad/runtime/ops_torch.py
@@ -6,7 +6,7 @@
 from tinygrad.runtime.lib import RawBuffer
 
 device = torch.device("cuda:0" if torch.cuda.is_available() else ("mps" if getenv("MPS", 0) else "cpu"))
-type_map = {torch.float16: dtypes.float16, torch.float32: dtypes.float32, torch.int8: dtypes.int8, torch.int32: dtypes.int32, torch.int64: dtypes.int64, torch.uint8: dtypes.uint8, torch.bool: dtypes.bool}
+type_map = {torch.float16: dtypes.float16, torch.float32: dtypes.float32, torch.float64: dtypes.float64, torch.double: dtypes.float64, torch.int8: dtypes.int8, torch.int32: dtypes.int32, torch.int64: dtypes.int64, torch.uint8: dtypes.uint8, torch.bool: dtypes.bool}
 inverse_type_map = {v:k for k,v in type_map.items()}
 
 torch_fxn_for_op: Dict[Op, Callable] = {**base_fxn_for_op, **{
diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py
index d5100a1fcd84b..dd557a251bbef 100644
--- a/tinygrad/tensor.py
+++ b/tinygrad/tensor.py
@@ -36,6 +36,7 @@ class Tensor:
 
   def __init__(self, data:Union[int, float, list, tuple, LazyBuffer, np.ndarray], device=Device.DEFAULT, dtype:Optional[DType]=None, requires_grad:Optional[bool]=None):
     assert dtype is None or isinstance(dtype, DType), f"invalid dtype {dtype}"
+    assert dtype is None or device not in dtype.excluded, f"dtype: {dtype} not supported on device: {device}"
     device = Device.canonicalize(device)
     if isinstance(data, (list, tuple)):
       data = np.array(data, dtype=(dtype if dtype is not None else Tensor.default_type).np)


Code review comments:
diff_hunk: @@ -36,6 +36,7 @@ class Tensor:
 
   def __init__(self, data:Union[int, float, list, tuple, LazyBuffer, np.ndarray], device=Device.DEFAULT, dtype:Optional[DType]=None, requires_grad:Optional[bool]=None):
     assert dtype is None or isinstance(dtype, DType), f"invalid dtype {dtype}"
+    assert dtype is None or device not in dtype.excluded, f"dtype: {dtype} not supported on device: {device}"
path: tinygrad/tensor.py
text: @user1:
You want this after the device `canonicalize` and you only want to check before the colon.

"gpu:0", "GPU:0" can both be passed into device

diff_hunk: @@ -87,7 +86,7 @@ def __call__(self, global_size, local_size, *bufs, wait=False) -> Optional[float
 
 class CLCodegen(CStyleCodegen):
   lang = CStyleLanguage(
-    kernel_prefix = "#define int64 long\n__kernel", buffer_prefix = "__global ", smem_prefix = "__local ",
+    kernel_prefix = "#ifdef cl_khr_fp64\n#pragma OPENCL EXTENSION cl_khr_fp64 : enable\n#elif defined(cl_amd_fp64)\n#pragma OPENCL EXTENSION cl_amd_fp64 : enable\n#endif\n#define int64 long\n__kernel", buffer_prefix = "__global ", smem_prefix = "__local ",
path: tinygrad/runtime/ops_gpu.py
text: @user1:
Want to do something similar to half_prekernel so this isn't in all the kernels?

@author:
misread this, ignore the last commit

diff_hunk: @@ -69,14 +72,15 @@ class dtypes:
   @staticmethod # static methds on top, or bool in the type info will refer to dtypes.bool
   def is_int(x: DType)-> bool: return x in (dtypes.int8, dtypes.uint8, dtypes.int32, dtypes.int64)
   @staticmethod
-  def is_float(x: DType) -> bool: return x in (dtypes.float16, dtypes.float32)
+  def is_float(x: DType) -> bool: return x in (dtypes.float16, dtypes.float32, dtypes.float64)
   @staticmethod
   def is_unsigned(x: DType) -> bool: return x in (dtypes.uint8)
   @staticmethod
   def from_np(x) -> DType: return asdict(dtypes())[np.dtype(x).name]
   bool: Final[DType] = DType(0, 1, "bool", bool)
   float16: Final[DType] = DType(0, 2, "half", np.float16)
   float32: Final[DType] = DType(4, 4, "float", np.float32)
+  float64: Final[DType] = DType(5, 8, "double", np.float64, ("METAL", "GPU" if OSX else ""))
path: tinygrad/helpers.py
text: @user1:
I don't think you want "" here if not OSX, I think you just don't want that element in the tuple

@author:
yeah will change to `("METAL", "GPU") if OSX else ()`