diff --git a/tinygrad/function.py b/tinygrad/function.py index 0f71b90923..c7823b223c 100644 --- a/tinygrad/function.py +++ b/tinygrad/function.py @@ -17,7 +17,9 @@ pm_transform_unique_const = PatternMatcher([ ]) pm_ctx = PatternMatcher([ - (UPat((Ops.BUFFER, Ops.BIND), name="x"), add_to_ctx), + (UPat(Ops.BUFFER, src=(UPat(Ops.UNIQUE), UPat(Ops.DEVICE)), name="x"), + lambda ctx,x: x.replace(src=(UOp(Ops.LUNIQUE, arg=next(ctx[1])), x.src[1])) if x.src[0].arg > ctx[2] else add_to_ctx(ctx,x)), + (UPat(Ops.BIND, name="x"), add_to_ctx), (UPat((Ops.AFTER, Ops.CONTIGUOUS), name="x"), lambda ctx,x: add_to_ctx(ctx,x) if not x.op_in_backward_slice_with_self(Ops.PARAM) and x.op_in_backward_slice_with_self(Ops.BUFFER) else None), ])+pm_transform_unique_const @@ -46,6 +48,7 @@ class _function(Generic[ReturnType]): # run it and do surgery later with Context(ALLOW_DEVICE_USAGE=getenv("DEVICE_IN_FUNCTION_BUG", 0)): _function.depth += 1 + unique_start = next(UOp.unique_num) ret = self.fxn(*args, **kwargs) _function.depth -= 1 if isinstance(ret, Tensor): @@ -65,7 +68,7 @@ class _function(Generic[ReturnType]): # the BUFFERs that are left are the implicit inputs num_explicit = len(call_uops) - uret = graph_rewrite(uret, pm_ctx, (call_uops, itertools.count(0)), bottom_up=True, name="get_implicit_inputs") + uret = graph_rewrite(uret, pm_ctx, (call_uops, itertools.count(0), unique_start), bottom_up=True, name="get_implicit_inputs") name = getattr(self.fxn, '__qualname__', None) or type(self.fxn).__qualname__ if not self.allow_implicit: implicit_buffers = [x for x in call_uops[num_explicit:] if x.op is Ops.BUFFER] diff --git a/tinygrad/mixin/__init__.py b/tinygrad/mixin/__init__.py index d24a7d9d60..07b9123178 100644 --- a/tinygrad/mixin/__init__.py +++ b/tinygrad/mixin/__init__.py @@ -7,7 +7,7 @@ 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.dtype import ConstType, DType, DTypeLike, 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 if TYPE_CHECKING: @@ -20,6 +20,8 @@ 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 empty(*shape, **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 @@ -46,13 +48,11 @@ class OpMixin(ElementwiseMixin, ReduceMixin): @classmethod def invalids(cls, *shape, **kwargs) -> Self: """ - Creates a tensor with the given shape, filled with Invalid. + Creates an anonymous uninitialized buffer with the given shape. - This is an alternative to Tensor.empty when you want an "anonymous" buffer. - - Eventually Tensor.empty will be replaced by this. + You can pass in `dtype` and `device` keyword arguments to control the data type and device of the tensor. """ - return cls.full(argfix(*shape), Invalid, **kwargs) + return cls.empty(argfix(*shape), **kwargs) @classmethod def zeros(cls, *shape, **kwargs) -> Self: