diff --git a/test/test_schedule.py b/test/test_schedule.py index d71e983578..6e2a999210 100644 --- a/test/test_schedule.py +++ b/test/test_schedule.py @@ -9,7 +9,7 @@ import functools from typing import List, Optional, Union, cast from tinygrad import nn, dtypes, Device, Tensor -from tinygrad.device import is_dtype_supported +from tinygrad.device import is_dtype_supported, view_supported_devices from tinygrad.dtype import DType from tinygrad.shape.shapetracker import ShapeTracker from tinygrad.shape.view import View @@ -18,7 +18,7 @@ from tinygrad.helpers import CI, DEBUG, FUSE_ARANGE, GlobalCounters, flatten, ge from tinygrad.codegen.kernel import Kernel, verify_ast from tinygrad.engine.schedule import BUF_LIMIT, ScheduleItem, create_schedule, view_right, view_left, do_realize from tinygrad.engine.realize import CompiledRunner, get_runner, run_schedule -from tinygrad.engine.lazy import LazyBuffer, view_supported_devices +from tinygrad.engine.lazy import LazyBuffer from extra.models.llama import precompute_freqs_cis class KernelCountException(Exception): pass diff --git a/test/test_subbuffer.py b/test/test_subbuffer.py index 18702d53d4..9d80820140 100644 --- a/test/test_subbuffer.py +++ b/test/test_subbuffer.py @@ -1,7 +1,6 @@ import unittest from tinygrad import Device, dtypes, Tensor -from tinygrad.device import Buffer -from tinygrad.engine.lazy import view_supported_devices +from tinygrad.device import Buffer, view_supported_devices @unittest.skipIf(Device.DEFAULT not in view_supported_devices, "subbuffer not supported") class TestSubBuffer(unittest.TestCase): diff --git a/tinygrad/device.py b/tinygrad/device.py index 98b0aaa821..8dae4e1220 100644 --- a/tinygrad/device.py +++ b/tinygrad/device.py @@ -43,6 +43,7 @@ Device = _Device() # **************** Buffer + Allocators **************** + @dataclass(frozen=True, eq=True) class BufferSpec: # TODO: move device, size, dtype here? @@ -53,6 +54,8 @@ class BufferSpec: nolru: bool = False external_ptr: Optional[int] = None +view_supported_devices = {"LLVM", "CLANG", "CUDA", "NV", "AMD", "METAL", "QCOM", "DSP", "DISK"} + class Buffer: def __init__(self, device:str, size:int, dtype:DType, opaque:Any=None, options:Optional[BufferSpec]=None, initial_value:Optional[bytes]=None, lb_refcount=0, base:Optional[Buffer]=None, offset:int=0, preallocate=False): diff --git a/tinygrad/engine/lazy.py b/tinygrad/engine/lazy.py index 74bd41a246..34943f0e72 100644 --- a/tinygrad/engine/lazy.py +++ b/tinygrad/engine/lazy.py @@ -5,7 +5,7 @@ from tinygrad.helpers import prod, getenv, all_int, all_same, DEBUG, _METADATA, from tinygrad.ops import exec_alu, python_alu from tinygrad.ops import identity_element, MathTrait, resolve, UOp, sint, GroupOp, Ops from tinygrad.shape.shapetracker import ShapeTracker -from tinygrad.device import Buffer +from tinygrad.device import Buffer, view_supported_devices from weakref import ref, ReferenceType, WeakValueDictionary lazycache: WeakValueDictionary[Any, LazyBuffer] = WeakValueDictionary() @@ -21,7 +21,6 @@ def create_lazybuffer(device:str, st:ShapeTracker, dtype:DType, op:Optional[Ops] if enable_cache: lazycache[cache_key] = ret return ret -view_supported_devices = {"LLVM", "CLANG", "CUDA", "NV", "AMD", "METAL", "QCOM", "DSP", "DISK"} class LazyBuffer(MathTrait): def __init__(self, device:str, st:ShapeTracker, dtype:DType, op:Optional[Ops]=None, arg:Any=None, srcs:Tuple[LazyBuffer, ...]=(),