nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

full(buffer=True) (#16311)

Browse Source 
make full a buffer with flag to turn off
This commit is contained in:
chenyu
2026-05-21 16:34:44 -04:00
committed by GitHub
parent 6815f28849
commit 73ea36f4ac
5 changed files with 26 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/null/test_tensor.py
-6
View File
@@ -145,12 +145,6 @@ class TestTensorUnique(unittest.TestCase):
Tensor.realize(a,b)
self.assertIsNot(a.uop.buffer, b.uop.buffer)
def test_eye_bufs_unique(self):
a = Tensor.eye(10).contiguous()
b = Tensor.eye(10).contiguous()
Tensor.realize(a,b)
self.assertIsNot(a.uop.buffer, b.uop.buffer)
def test_times_2_not_unique(self):
a = Tensor.zeros(10, 10).contiguous()
b = a * 2
test/unit/test_setitem_schedule.py
+3 -1
View File
@@ -81,7 +81,7 @@ class TestSetitemInto(unittest.TestCase):
self.assertEqual(GlobalCounters.kernel_count, 1)
self.assertListEqual(t.tolist(), [2, 5, 4, 5])
def test_setitem_into_cont(self):
def test_setitem_into_const(self):
GlobalCounters.reset()
t = Tensor.ones(4, dtype=dtypes.int32)
t[1] = 5
@@ -110,7 +110,9 @@ class TestSetitemInto(unittest.TestCase):
def test_setitem_into_arange(self):
# NOTE: arange has no real buffer, but assigning to it is fine
GlobalCounters.reset()
other = Tensor.arange(4, dtype=dtypes.int32)
t = Tensor.arange(4, dtype=dtypes.int32)
self.assertIs(other.uop, t.uop)
t[1] = 5
self.assertEqual(GlobalCounters.kernel_count, 0)
t.realize()
tinygrad/mixin/__init__.py
+14 -6
View File
@@ -6,6 +6,7 @@ from tinygrad.mixin.movement import MovementMixin
from tinygrad.mixin.reduce import ReduceMixin
from tinygrad.uop import Ops
from tinygrad.uop.ops import _broadcast_shape, resolve, smax, smin, identity_element
from tinygrad.device import canonicalize_device
from tinygrad.dtype import ConstType, DType, DTypeLike, Invalid, InvalidType, PtrDType, PyConst, dtypes, least_upper_dtype, sum_acc_dtype, to_dtype
from tinygrad.helpers import all_int, argfix, ceildiv, flatten, flat_to_grouped, make_tuple, prod, resolve_pool_pads, round_up
@@ -18,6 +19,8 @@ ReductionStr = Literal["mean", "sum", "none"]
class OpMixin(ElementwiseMixin, ReduceMixin):
@staticmethod
def unique_const(fill_value:ConstType, **kwargs): raise NotImplementedError("creation helpers are only supported on Tensor and UOp")
@staticmethod
def const(dtype, b, device=None): raise NotImplementedError("creation helpers are only supported on Tensor and UOp")
@classmethod
def full(cls, shape:tuple[sint, ...], fill_value:ConstType, **kwargs) -> Self:
@@ -34,7 +37,11 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
print(Tensor.full((2, 3), False).numpy())
```
"""
return cls.unique_const(fill_value, **kwargs).reshape((1,)*len(new_shape := argfix(shape))).expand(new_shape)
new_shape = argfix(shape)
if not kwargs.pop("buffer", True):
dt = to_dtype(kwargs.pop("dtype", None) or dtypes.from_py(fill_value))
return cls.const(dt, fill_value, canonicalize_device(kwargs.pop("device", None))).reshape((1,)*len(new_shape)).expand(new_shape)
return cls.unique_const(fill_value, **kwargs).reshape((1,)*len(new_shape)).expand(new_shape)
@classmethod
def invalids(cls, *shape, **kwargs) -> Self:
@@ -112,7 +119,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
if lo < (dt:=to_dtype(dtype)).min or dt.max < hi: raise OverflowError(f"arange [{start}, {stop}) is not representable in dtype {dtype}")
# NOTE: this matches numpy, torch raises RuntimeError if stop-start and step have different signs
if (output_len:=ceildiv(stop-start, step)) <= 0: return cls.full((0,), 0, dtype=dtype, **kwargs)
return (cls.full((output_len,), step, dtype=dtype, **kwargs)._cumalu(0, Ops.ADD) + (start - step)).cast(dtype)
return (cls.full((output_len,), step, dtype=dtype, buffer=False, **kwargs)._cumalu(0, Ops.ADD) + (start - step)).cast(dtype)
@classmethod
def linspace(cls, start:int|float, stop:int|float, steps:int, **kwargs) -> Self:
@@ -711,7 +718,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
if self.ndim == 0: return self._split_cumalu(axis, Ops.MAX), type(self).zeros(self.shape, dtype=dtypes.int32, device=self.device)
values, n = self._split_cumalu(axis, Ops.MAX), int(self.shape[axis])
x, values_t = self.transpose(axis, -1), values.transpose(axis, -1)
match = x.unsqueeze(-1).eq(values_t.unsqueeze(-2)) * type(self).ones(n, n, device=self.device).triu()
match = x.unsqueeze(-1).eq(values_t.unsqueeze(-2)) * type(self).ones(n, n, device=self.device, buffer=False).triu()
idx = (-(match * type(self).arange(n, 0, -1, device=self.device).reshape(n, 1)).max(-2) + n).cast(dtypes.int32)
return values, idx.transpose(-1, axis)
@@ -758,7 +765,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
last_dim_size = x.shape[-1]
x_unsqueezed = x.unsqueeze(-2).expand((None,)*(self.ndim-1)+(last_dim_size, None))
x_cummax, _ = x.cummax(-1)
mask = type(self).ones(last_dim_size, last_dim_size, device=self.device).tril()
mask = type(self).ones(last_dim_size, last_dim_size, device=self.device, buffer=False).tril()
ret = mask.where(x_unsqueezed - x_cummax.unsqueeze(-1), self.dtype.min).exp().sum(-1).log() + x_cummax
return ret.transpose(-1, axis)
@@ -855,7 +862,8 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
x = blue_box.cat(flipped_green_box.flip(flip_dims), dim=crossover_dim)
x = x.flatten(dim, dim+n_stages-1).shrink_to(self.shape)
# compute indices for sorted values
mask = type(self).ones(orig_len, orig_len, dtype=dtypes.bool, device=self.device).tril().reshape((None, None) + (1,)*(self.ndim-dim-1))
mask = type(self).ones(orig_len, orig_len, dtype=dtypes.bool, device=self.device, buffer=False).tril()
mask = mask.reshape((None, None) + (1,)*(self.ndim-dim-1))
def compute_counts(t:Self): return (mask & t.unsqueeze(dim).eq(t.unsqueeze(dim+1))).sum(dim+1)
count_orig, count_sorted = compute_counts(self), compute_counts(x)
cond = self.unsqueeze(dim+1).eq(x.unsqueeze(dim)) & count_orig.unsqueeze(dim+1).eq(count_sorted.unsqueeze(dim))
@@ -1061,7 +1069,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
```
"""
if reduce not in {None, "add", "multiply"}: raise TypeError(f"{reduce=} must be one of None, 'multiply', or 'add'")
if isinstance(src, (int, float, bool)): src = type(self).full(index.shape, src, dtype=self.dtype, device=self.device)
if isinstance(src, (int, float, bool)): src = type(self).full(index.shape, src, dtype=self.dtype, device=self.device, buffer=False)
elif reduce: raise TypeError("non-scalar src is not supported with reduce arg. use scatter_reduce")
if reduce == "add": return self.scatter_reduce(dim, index, src, "sum", include_self=True)
if reduce == "multiply": return self.scatter_reduce(dim, index, src, "prod", include_self=True)
tinygrad/tensor.py
+8 -4
View File
@@ -160,6 +160,9 @@ class Tensor(OpMixin):
def alu(self, op: Ops, *src: Tensor) -> Tensor: return self._apply_uop(lambda *u: u[0].alu(op, *u[1:]), *src)
def const_like(self, b:ConstType) -> Tensor: return Tensor(self.uop.const_like(b), requires_grad=False)
@staticmethod
def const(dtype:DType, b:ConstType|UOp, device:str|tuple[str, ...]|None=None) -> Tensor:
return Tensor(b if isinstance(b, UOp) else UOp.const(dtype, b, device))
@staticmethod
def unique_const(fill_value:ConstType|UOp, **kwargs) -> Tensor:
if isinstance(fill_value, UOp): return Tensor(fill_value, **kwargs)
dtype, device = kwargs.pop("dtype", None), kwargs.pop("device", None)
@@ -1008,10 +1011,11 @@ class Tensor(OpMixin):
if isinstance(v, Tensor) and v.dtype != self.dtype: raise RuntimeError(f"setitem dtype mismatch: {self.dtype=} != {v.dtype=}")
# raise if mutation would diverge from eager (allow only pure views of a realized buffer; exclude +=/-= RHS via v_uop/v_bw)
v_uop, v_bw = (v.uop, v.uop.backward_slice) if isinstance(v, Tensor) else (None, {})
shared = self.uop.base if self.uop.base.is_realized else None
if any(self.uop in t.uop.backward_slice_with_self and t.uop.base is not shared for tref in all_tensors
if (t:=tref()) is not None and t is not self and t.uop is not v_uop and t.uop not in v_bw):
raise RuntimeError("can't setitem on a tensor with other uses")
if self.uop.op_in_backward_slice_with_self(Ops.BUFFER):
shared = self.uop.base if self.uop.base.is_realized else None
if any(self.uop in t.uop.backward_slice_with_self and t.uop.base is not shared for tref in all_tensors
if (t:=tref()) is not None and t is not self and t.uop is not v_uop and t.uop not in v_bw):
raise RuntimeError("can't setitem on a tensor with other uses")
if not self.uop.base.is_realized and self.is_floating_point() and (self.requires_grad or (isinstance(v, Tensor) and v.requires_grad)):
if not isinstance(v, Tensor): v = Tensor(v, device=self.device, dtype=self.dtype)
# __iadd__/__isub__ creates AFTER(view, STORE(view, computed)); unwrap to get the computed value
tinygrad/uop/ops.py
+1 -1
View File
@@ -1208,7 +1208,7 @@ class UPat(OpMixin):
def cvar(name:str|None=None, dtype:DType|tuple[DType, ...]|None=None, vec=True, arg=None):
return UPat(Ops.CONST, dtype, name=name, arg=arg)
@staticmethod
def const(dtype:DType|tuple[DType, ...]|None, b:ConstType): return UPat(Ops.CONST, dtype=dtype, arg=b)
def const(dtype:DType|tuple[DType, ...]|None, b:ConstType, device=None): return UPat(Ops.CONST, dtype=dtype, arg=b)
# lil helper
def f(self, op, **kwargs): return UPat(op, src=(self,), **kwargs)