tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 18:36:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

cleanup Device usage in Tensor (#12662)

Browse Source
This commit is contained in:
George Hotz
2025-10-14 16:22:22 +08:00
committed by GitHub
parent e06cbfcb8a
commit 30ee7c4c26
1 changed files with 27 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tinygrad/tensor.py
+27 -23
View File
@@ -19,6 +19,9 @@ from tinygrad.engine.schedule import ScheduleItem, create_schedule_with_vars
from tinygrad.schedule.rangeify import get_rangeify_map
from tinygrad.schedule.multi import get_multi_map
# TODO: this should be the only usage of Device
def canonicalize_device(device:str|None) -> str: return Device.canonicalize(device)
# *** all in scope Tensors are here. this gets relevant UOps ***
all_tensors: dict[weakref.ref[Tensor], None] = {}
@@ -113,9 +116,10 @@ class Tensor(MathTrait):
def __init__(self, data:ConstType|bytes|list|tuple|UOp|'np.ndarray'|pathlib.Path|None, # type: ignore [name-defined] # noqa: F821
device:str|tuple|list|None=None, dtype:DTypeLike|None=None, requires_grad:bool|None=None):
if dtype is not None: dtype = to_dtype(dtype)
if device is None and isinstance(data, pathlib.Path): device = f"DISK:{data.resolve()}" # keep it on the disk if device is None
device = tuple(Device.canonicalize(x) for x in device) if isinstance(device, (tuple, list)) else Device.canonicalize(device)
_dtype:DType|None = to_dtype(dtype) if dtype is not None else None
_device:str|tuple[str, ...] = tuple(canonicalize_device(x) for x in device) if isinstance(device, (tuple, list)) else canonicalize_device(device)
del device, dtype
# tensors can have gradients if you have called .backward
self.grad:Tensor|None = None
@@ -126,41 +130,41 @@ class Tensor(MathTrait):
# create a UOp from the different types of inputs
if isinstance(data, UOp):
assert dtype is None or dtype==data.dtype, "dtype doesn't match, and casting isn't supported"
assert _dtype is None or _dtype==data.dtype, "dtype doesn't match, and casting isn't supported"
# if data is dtype.index that means that this is a symbolic int and we need to lower it to something we can make a Tensor out of
if data.dtype==dtypes.index: data = _index_to_concrete_int(data)
if data.op is Ops.BIND: # type: ignore # mypy type narrowing is bugged here
var, val = data.unbind() # type: ignore
# give the bound constant a device
const = UOp.const(var.dtype, val, device, ())
const = UOp.const(var.dtype, val, _device, ())
data = data.replace(src=(var.replace(src=const.src), const)) # type: ignore
elif data is None: data = UOp.const(dtype or dtypes.default_float, 0, device, ())
elif isinstance(data, get_args(ConstType)): data = UOp.const(dtype or dtypes.from_py(data), data, device, ())
elif isinstance(data, bytes): data = _frompy(data, dtypes.uint8 if dtype is None else dtype)
elif data is None: data = UOp.const(_dtype or dtypes.default_float, 0, _device, ())
elif isinstance(data, get_args(ConstType)): data = UOp.const(_dtype or dtypes.from_py(data), data, _device, ())
elif isinstance(data, bytes): data = _frompy(data, dtypes.uint8 if _dtype is None else _dtype)
elif isinstance(data, (list, tuple)):
if dtype is None:
if (d := fully_flatten(data)) and all(isinstance(s, bool) for s in d): dtype = dtypes.bool
else: dtype = dtypes.default_int if d and all_int(d) else dtypes.default_float # NOTE: this works because all_int([True, False]) is True
if dtype in [dtypes.bfloat16, *dtypes.fp8s]: data = Tensor(_frompy(data, dtypes.float32), device=device).cast(dtype).uop
else: data = _frompy(data, dtype)
if _dtype is None:
if (d := fully_flatten(data)) and all(isinstance(s, bool) for s in d): _dtype = dtypes.bool
else: _dtype = dtypes.default_int if d and all_int(d) else dtypes.default_float # NOTE: this works because all_int([True, False]) is True
if _dtype in [dtypes.bfloat16, *dtypes.fp8s]: data = Tensor(_frompy(data, dtypes.float32), device=_device).cast(_dtype).uop
else: data = _frompy(data, _dtype)
elif is_numpy_ndarray(data):
import numpy as np
assert isinstance(data, np.ndarray), f"expected np.ndarray, got {data}"
if data.shape == (): data = UOp.const(dtype or _from_np_dtype(data.dtype), data.item(), device, ())
else: data = _fromnp(data.astype(npdtype) if dtype is not None and (npdtype:=_to_np_dtype(dtype)) is not None else data) # type: ignore [name-defined]
if data.shape == (): data = UOp.const(_dtype or _from_np_dtype(data.dtype), data.item(), _device, ())
else: data = _fromnp(data.astype(npdtype) if _dtype is not None and (npdtype:=_to_np_dtype(_dtype)) is not None else data) # type: ignore [name-defined]
elif isinstance(data, pathlib.Path):
dtype = dtype or dtypes.uint8
data = UOp.new_buffer(f"DISK:{data.resolve()}", data.stat().st_size // dtype.itemsize, dtype)
_dtype = _dtype or dtypes.uint8
data = UOp.new_buffer(f"DISK:{data.resolve()}", data.stat().st_size // _dtype.itemsize, _dtype)
# by this point, it has to be a UOp
if not isinstance(data, UOp): raise RuntimeError(f"can't create Tensor from {data!r} with type {type(data)}")
# data might be on a different device
if isinstance(device, str): self.uop:UOp = data if data.device == device else data.copy_to_device(device)
if isinstance(_device, str): self.uop:UOp = data if data.device == _device else data.copy_to_device(_device)
# if device is a tuple, we should have/construct a MultiLazyBuffer
elif isinstance(data.device, str): self.uop = Tensor(data).shard(device).uop
elif isinstance(data.device, str): self.uop = Tensor(data).shard(_device).uop
else:
assert data.device == device, f"MultiLazyBuffer device mismatch, {data.device} != {device}"
assert data.device == _device, f"MultiLazyBuffer device mismatch, {data.device} != {_device}"
self.uop = data
# add to all_tensors after construction succeeds
@@ -376,7 +380,7 @@ class Tensor(MathTrait):
"""
Moves the tensor to the given device.
"""
device = tuple(Device.canonicalize(x) for x in device) if isinstance(device, (tuple, list)) else Device.canonicalize(device)
device = tuple(canonicalize_device(x) for x in device) if isinstance(device, (tuple, list)) else canonicalize_device(device)
if device == self.device: return self
if not isinstance(device, str): return self.shard(device)
ret = Tensor(self.uop, device, requires_grad=self.requires_grad)
@@ -401,7 +405,7 @@ class Tensor(MathTrait):
```
"""
assert isinstance(self.device, str), "can't shard a MultiLazyBuffer"
devices = tuple(Device.canonicalize(x) for x in devices)
devices = tuple(canonicalize_device(x) for x in devices)
mlb = self.uop.shard(devices, self._resolve_dim(axis)) if axis is not None else self.uop.copy_to_device(devices)
return Tensor(mlb, device=devices, requires_grad=self.requires_grad)
@@ -490,7 +494,7 @@ class Tensor(MathTrait):
dtype, shape = to_dtype(dtype) if dtype is not None else dtypes.default_float, argfix(*shape)
if not isinstance(size:=prod([x.vmax if isinstance(x, UOp) else x for x in shape]), int): raise ValueError(f"size must be int {size}")
# TODO: add test for multidevice tensor
device = tuple(Device.canonicalize(d) for d in device) if isinstance(device, tuple) else Device.canonicalize(device)
device = tuple(canonicalize_device(d) for d in device) if isinstance(device, tuple) else canonicalize_device(device)
return Tensor(UOp.new_buffer(device, size, dtype), device, dtype, **kwargs).shrink(((0,prod(shape)),)).reshape(shape)
def empty_like(self, **kwargs) -> Tensor:
@@ -572,7 +576,7 @@ class Tensor(MathTrait):
if not dtypes.is_float(dtype := to_dtype(dtype or dtypes.default_float)): raise ValueError(f"rand only supports float dtypes, got {dtype}")
if not all_int(shape:=argfix(*shape)) or not all(s >= 0 for s in shape): raise ValueError(f"invalid input {shape=}")
if device is not None and not isinstance(device, str): raise ValueError(f"rand only supports single device, got {device=}")
device = Device.canonicalize(device)
device = canonicalize_device(device)
# if shape has 0, return zero tensor
if (numel := prod(shape)) == 0: return Tensor.zeros(shape, device=device, dtype=dtype, **kwargs)