From 73ea36f4acb38e24c127c2b5fce7f94fd4db3319 Mon Sep 17 00:00:00 2001 From: chenyu Date: Thu, 21 May 2026 16:34:44 -0400 Subject: [PATCH] full(buffer=True) (#16311) make full a buffer with flag to turn off --- test/null/test_tensor.py | 6 ------ test/unit/test_setitem_schedule.py | 4 +++- tinygrad/mixin/__init__.py | 20 ++++++++++++++------ tinygrad/tensor.py | 12 ++++++++---- tinygrad/uop/ops.py | 2 +- 5 files changed, 26 insertions(+), 18 deletions(-) diff --git a/test/null/test_tensor.py b/test/null/test_tensor.py index db65a5692b..7aee67926c 100644 --- a/test/null/test_tensor.py +++ b/test/null/test_tensor.py @@ -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 diff --git a/test/unit/test_setitem_schedule.py b/test/unit/test_setitem_schedule.py index 7527fd315d..044932bae2 100644 --- a/test/unit/test_setitem_schedule.py +++ b/test/unit/test_setitem_schedule.py @@ -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() diff --git a/tinygrad/mixin/__init__.py b/tinygrad/mixin/__init__.py index 3984627cbe..487a8a1395 100644 --- a/tinygrad/mixin/__init__.py +++ b/tinygrad/mixin/__init__.py @@ -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) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index fbdd120530..60e06f6482 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -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 diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 94b0034a45..82bf5ebe25 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -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)