mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-08-29 12:16:08 +00:00
Device._buffers -> Device._devices (#3052)
backend devices used to be called buffers
This commit is contained in:
+1
-1
@@ -6,7 +6,7 @@ from tinygrad.helpers import getenv, CI
|
||||
def multidevice_test(fxn):
|
||||
exclude_devices = getenv("EXCLUDE_DEVICES", "").split(",")
|
||||
def ret(self):
|
||||
for device in Device._buffers:
|
||||
for device in Device._devices:
|
||||
if device in ["DISK", "FAKE"]: continue
|
||||
if not CI: print(device)
|
||||
if device in exclude_devices:
|
||||
|
||||
+3
-3
@@ -15,18 +15,18 @@ if TYPE_CHECKING:
|
||||
# **************** Device ****************
|
||||
|
||||
class _Device:
|
||||
def __init__(self) -> None: self._buffers: List[str] = [x.stem[len("ops_"):].upper() for x in (pathlib.Path(__file__).parent/"runtime").iterdir() if x.stem.startswith("ops_")] # noqa: E501
|
||||
def __init__(self) -> None: self._devices: List[str] = [x.stem[len("ops_"):].upper() for x in (pathlib.Path(__file__).parent/"runtime").iterdir() if x.stem.startswith("ops_")] # noqa: E501
|
||||
def canonicalize(self, device:Optional[str]) -> str: return (device.split(":", 1)[0].upper() + ((":"+device.split(":", 1)[1]) if ':' in device else '')).replace(":0", "") if device is not None else self.DEFAULT # noqa: E501
|
||||
def __getitem__(self, ix:str) -> Union[Interpreted, Compiled]: return self.__get_canonicalized_item(self.canonicalize(ix))
|
||||
@functools.lru_cache(maxsize=None) # this class is a singleton, pylint: disable=method-cache-max-size-none
|
||||
def __get_canonicalized_item(self, ix:str) -> Union[Interpreted, Compiled]:
|
||||
x = ix.split(":")[0].upper()
|
||||
ret = [cls for cname, cls in inspect.getmembers(importlib.import_module(f'tinygrad.runtime.ops_{x.lower()}')) if (cname.lower() == x.lower() + "device") and x in self._buffers][0] # noqa: E501
|
||||
ret = [cls for cname, cls in inspect.getmembers(importlib.import_module(f'tinygrad.runtime.ops_{x.lower()}')) if (cname.lower() == x.lower() + "device") and x in self._devices][0] # noqa: E501
|
||||
if isinstance(ret, type): ret = ret(ix)
|
||||
return ret
|
||||
@functools.cached_property
|
||||
def DEFAULT(self) -> str:
|
||||
device_from_env: Optional[str] = functools.reduce(lambda val, ele: ele if getenv(ele) == 1 else val, self._buffers, None) # type: ignore
|
||||
device_from_env: Optional[str] = functools.reduce(lambda val, ele: ele if getenv(ele) == 1 else val, self._devices, None) # type: ignore
|
||||
if device_from_env: return device_from_env
|
||||
for device in ["METAL", "CUDA", "HIP", "GPU"]:
|
||||
try:
|
||||
|
||||
+1
-1
@@ -936,7 +936,7 @@ class Tensor:
|
||||
def is_floating_point(self) -> bool: return dtypes.is_float(self.dtype)
|
||||
|
||||
# register functions to move between devices
|
||||
for device in Device._buffers: setattr(Tensor, f"{device.lower()}", partialmethod(Tensor.to, device))
|
||||
for device in Device._devices: setattr(Tensor, f"{device.lower()}", partialmethod(Tensor.to, device))
|
||||
|
||||
if IMAGE:
|
||||
# if IMAGE>0 we install these replacement functions in Tensor (hack!)
|
||||
|
||||
Reference in New Issue
Block a user