diff --git a/test/unit/test_call.py b/test/unit/test_call.py index 0dd805f54e..5fcb690fd6 100644 --- a/test/unit/test_call.py +++ b/test/unit/test_call.py @@ -3,6 +3,9 @@ import numpy as np from tinygrad import Tensor, function, Device from tinygrad.dtype import dtypes from tinygrad.uop.ops import UOp, Ops +from tinygrad.tensor import transform_to_call + +def sched_key(t:Tensor): return transform_to_call(UOp.sink(t.uop))[0].src[0].key class TestCall(unittest.TestCase): def test_call_plus(self): @@ -223,9 +226,7 @@ class TestCallSchedule(unittest.TestCase): a = Tensor.ones(3) x = f(a, UOp.variable("scale_a", 1, 100).bind(2)) y = f(a, UOp.variable("scale_b", 1, 100).bind(3)) - fx = next(u for u in x.uop.toposort() if u.op is Ops.CALL and u.num_returned) - fy = next(u for u in y.uop.toposort() if u.op is Ops.CALL and u.num_returned) - self.assertEqual(fx.src[0].key, fy.src[0].key) + self.assertEqual(sched_key(x), sched_key(y)) np.testing.assert_equal(x.numpy(), [2, 2, 2]) np.testing.assert_equal(y.numpy(), [3, 3, 3]) @@ -245,17 +246,17 @@ class TestCallSchedule(unittest.TestCase): np.testing.assert_equal(cache.numpy()[8:], np.zeros(8)) def test_precompile_schedule_cache_hit(self): - """two instances of the same @function should produce identical function body keys (schedule cache hit)""" + """two instances of the same @function should produce identical scheduled function keys without aliasing their outputs""" @function(precompile=True) def f(x:Tensor) -> Tensor: return x + Tensor.full(x.shape, -1.0) a = Tensor.empty(4, 8) b = Tensor.empty(4, 8) r0, r1 = f(a), f(b) - # find the call nodes c0 = next(u for u in r0.uop.toposort() if u.op is Ops.CALL and u.num_returned) c1 = next(u for u in r1.uop.toposort() if u.op is Ops.CALL and u.num_returned) - # the function bodies (src[0]) should have identical keys - self.assertEqual(c0.src[0].key, c1.src[0].key) + # output identities stay unique per call; they canonicalize only when combined into a scheduling scope + self.assertIsNot(c0.src[-1], c1.src[-1]) + self.assertEqual(sched_key(r0), sched_key(r1)) def test_precompile_consumes_call_output(self): """a precompiled function consuming the output of a non-precompiled function""" @@ -288,9 +289,9 @@ class TestCallSchedule(unittest.TestCase): class TestArgOrder(unittest.TestCase): """RETURNED placeholders can appear anywhere in a call's srcs: slots are src positions, nothing reorders""" def make_intersperse_call(self, x, precompile=False): - # call with sources (body, returned(slot=0), input(slot=1)): the input is the input, the output binds the RETURNED + # call with sources (body, returned, input(slot=1)): the input is the input, the output binds the RETURNED dev = x.device if isinstance(x.device, str) else (x.device or (Device.DEFAULT,))[0] - r0 = UOp.returned(0, x.dtype, x.shape, device=dev) + r0 = UOp.returned(x.dtype, x.shape, device=dev) o0 = UOp.param(0, x.dtype, x.shape, dev) p1 = UOp.param(1, x.dtype, x.shape, dev) from tinygrad.uop.ops import CallInfo @@ -323,7 +324,7 @@ class TestArgOrder(unittest.TestCase): x = Tensor([1.0, 2.0, 3.0]).realize() x.requires_grad = True dev = x.device if isinstance(x.device, str) else (x.device or (Device.DEFAULT,))[0] - r0 = UOp.returned(0, dtypes.float, x.shape, device=dev) + r0 = UOp.returned(dtypes.float, x.shape, device=dev) o0 = UOp.param(0, dtypes.float, x.shape, dev) p1 = UOp.param(1, dtypes.float, x.shape, dev) from tinygrad.uop.ops import CallInfo diff --git a/tinygrad/function.py b/tinygrad/function.py index 4ba84d7c0b..fb0f8645d6 100644 --- a/tinygrad/function.py +++ b/tinygrad/function.py @@ -13,9 +13,10 @@ def add_to_ctx(ctx, x:UOp): return ret pm_ctx = PatternMatcher([ - (UPat(Ops.BUFFER, 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), + # unbound BUFFERs and their AFTER outputs are scoped inside their CALL: they are never implicit inputs + (UPat(Ops.BUFFER, name="x"), lambda ctx,x: None if x.is_unbound else add_to_ctx(ctx,x)), + (UPat((Ops.AFTER, Ops.CONTIGUOUS), name="x"), lambda ctx,x: add_to_ctx(ctx,x) if not x.buf_uop.is_unbound and + not x.op_in_backward_slice_with_self(Ops.PARAM) and x.op_in_backward_slice_with_self(Ops.BUFFER) else None), ]) def invalid_outputs(uret:UOp) -> set[UOp]: diff --git a/tinygrad/mixin/gradient.py b/tinygrad/mixin/gradient.py index 93e5d213bb..f0ec320e79 100644 --- a/tinygrad/mixin/gradient.py +++ b/tinygrad/mixin/gradient.py @@ -29,14 +29,14 @@ def call_gradient(ctx:UOp, k:UOp, needed:set[int]) -> tuple[UOp|None, ...]: # grads align with the call's src positions (None for the body and for RETURNED outputs, wherever they are) def arg_grads(g): git = iter(g) - return (None,) + tuple(next(git) if a.unsharded_base.op is not Ops.RETURNED else None for a in k.src[1:]) + return (None,) + tuple(next(git) if not a.unsharded_base.is_unbound else None for a in k.src[1:]) if ctx.op is Ops.SINK: real = [on_dev(g, i) for i,g in enumerate(ctx.src) if g.op is not Ops.NOOP] return arg_grads(k.arg.grad_fxn(*real, call=k) if len(real) > 1 else k.arg.grad_fxn(real[0], k)) return arg_grads(k.arg.grad_fxn(on_dev(ctx, 0), k)) # the RETURNED inputs are the call outputs: their positions in the args get the output gradients from the AFTER rule assert fxn.op is Ops.SINK and k.num_returned, f"expected a CALL with RETURNED inputs or a grad_fxn, got {fxn.op}" - ret_pos = [i for i, a in enumerate(args) if a.unsharded_base.op is Ops.RETURNED] + ret_pos = [i for i, a in enumerate(args) if a.unsharded_base.is_unbound] # the body stores the outputs into output PARAMs: the values are the stored values in slot order values = UOp.sink(*[st.src[1] for st in fxn.src if st.op is Ops.STORE]) params = {x.arg.slot:x for x in fxn.toposort(enter_calls=False) if x.op == Ops.PARAM} diff --git a/tinygrad/nn/__init__.py b/tinygrad/nn/__init__.py index aae4600cc2..949e68c5c8 100644 --- a/tinygrad/nn/__init__.py +++ b/tinygrad/nn/__init__.py @@ -305,7 +305,7 @@ class RMSNorm: from tinygrad.uop.ops import UOp, KernelInfo, Ops, AxisType def _embedding_bwd(grad_emb:UOp, call:UOp) -> tuple: - weight, idx = (a for a in call.src[1:] if a.unsharded_base.op is not Ops.RETURNED) + weight, idx = (a for a in call.src[1:] if not a.unsharded_base.is_unbound) is_vocab_sharded = isinstance(weight.device, tuple) and weight.axis == 0 # for multi-device: replicate grad_emb and idx on all devices if isinstance(weight.device, tuple): diff --git a/tinygrad/schedule/__init__.py b/tinygrad/schedule/__init__.py index b0543d4cb6..314e8404bc 100644 --- a/tinygrad/schedule/__init__.py +++ b/tinygrad/schedule/__init__.py @@ -117,9 +117,12 @@ pm_resolve_linear_call = PatternMatcher([ schedule_cache: dict[bytes, UOp] = {} # ctx is just for DEBUG on inner -def lower_sink_to_linear(function:UOp) -> UOp|None: +def lower_sink_to_linear(call:UOp) -> UOp|None: + function = call.src[0] + if function.op is not Ops.SINK or isinstance(function.arg, KernelInfo): return None + # value calls (with RETURNED outputs) are inlined positionally during prepare: their bodies are not programs to schedule + if any(x.unsharded_base.is_unbound for x in call.src[1:]): return None st = time.perf_counter() - if isinstance(function.arg, KernelInfo): return None cache_key = function.key if not SCACHE or (sc_ret:=schedule_cache.get(cache_key, None)) is None: if SPEC: type_verify(function, spec_tensor) @@ -139,10 +142,10 @@ def lower_sink_to_linear(function:UOp) -> UOp|None: print(f"scheduled {len(linear.src):5d} kernels in {(time.perf_counter()-st)*1000:8.2f} ms"+\ f" | {' cache hit' if SCACHE and sc_ret is not None else 'CACHE MISS'} {cache_key.hex()[:8]}"+\ f" | {len(UOpMetaClass.ucache):7d} uops in cache"+("" if frm is None else f" | {frm.filename}:{frm.lineno}")) - return linear + return call.replace(src=(linear,)+call.src[1:]) pm_schedule = PatternMatcher([ - (UPat(Ops.SINK, name="function"), lower_sink_to_linear), + (UPat(Ops.CALL, name="call"), lower_sink_to_linear), ]) def assert_all_same_devices(ast:UOp): diff --git a/tinygrad/schedule/prepare.py b/tinygrad/schedule/prepare.py index b45e9d43cc..500a249c08 100644 --- a/tinygrad/schedule/prepare.py +++ b/tinygrad/schedule/prepare.py @@ -1,7 +1,7 @@ import itertools from tinygrad.dtype import dtypes, to_dtype from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp -from tinygrad.uop.ops import graph_rewrite, rewrite_group, ParamArg, identity_element, resolve_returned_after +from tinygrad.uop.ops import graph_rewrite, rewrite_group, identity_element, resolve_returned_after from tinygrad.uop.movement import mop_cleanup from tinygrad.helpers import prod, getenv, all_int, DEBUG, SPLIT_REDUCEOP, OPENPILOT_HACKS, FLOAT16, argsort from tinygrad.schedule.indexing import apply_movement_op @@ -202,7 +202,7 @@ def convert_copy_to_store(ctx, copy:UOp, existing_buf:UOp|None=None): # if there's already a buffer, we just use it return existing_buf.flatten().store(input_src) # create the output buffer - buf = UOp(Ops.BUFFER, arg=ParamArg(next(ctx), copy.dtype, size=prod(input_src.max_shape), device=copy.device)) + buf = UOp.new_buffer(copy.device, prod(input_src.max_shape), copy.dtype) # reshape back to input return buf.reshape(input_src.max_shape).after(buf.store(input_src)).reshape(copy.shape) diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index d2bbdf2eb6..d89c1d28ce 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -229,7 +229,7 @@ def bufferize_to_store(ctx:itertools.count, x:UOp, idx:UOp, allow_locals=True): # NOTE: the local BUFFER needs to be disambiguated here if x.arg.addrspace == AddrSpace.GLOBAL: - buf = UOp(Ops.BUFFER, arg=ParamArg(next(ctx), dtype, size=size, device=x.arg.device, addrspace=AddrSpace.GLOBAL)) + buf = UOp.new_buffer(x.arg.device, size, dtype) do_store = buf.index(idx).store(x.src[0].cast(dtype)).end(*rngs) return buf.after(do_store).cast(x.dtype) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 9c1bdd0a72..ebcb40f03a 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -1,7 +1,7 @@ # inspired by https://github.com/karpathy/micrograd/blob/master/micrograd/engine.py from __future__ import annotations import time, functools, sys, inspect, pathlib, hashlib, weakref -from dataclasses import dataclass, field +from dataclasses import dataclass, field, replace from typing import Any, Callable, cast, get_args, ParamSpec, TypeGuard, TypeVar, Generic, TYPE_CHECKING if TYPE_CHECKING: import numpy from tinygrad.dtype import DType, DTypeLike, dtypes, ConstType, least_upper_dtype, to_dtype, _from_np_dtype, _to_np_dtype, PyConst, AddrSpace @@ -24,6 +24,7 @@ class AllocCtx: bases: set[UOp] = field(default_factory=set) stores: list[UOp] = field(default_factory=list) replacements: list[UOp] = field(default_factory=list) + unbound: dict[UOp, UOp] = field(default_factory=dict) views: set[UOp] = field(default_factory=set) # a tag is the tuple of original pre-rewrite UOps a node provides storage for @@ -94,7 +95,7 @@ def transform_precompiled_call(c:UOp) -> UOp|None: if c.arg is None or not c.arg.precompile or c.num_returned == 0: return None assert c.src[0].op is Ops.SINK, "precompiled call bodies are SINKs of stores into the output PARAMs" # the RETURNED srcs are the call outputs (slots are src positions) - ret_pos = [p for p,a in enumerate(c.src[1:]) if a.unsharded_base.op is Ops.RETURNED] + ret_pos = [p for p,a in enumerate(c.src[1:]) if a.unsharded_base.is_unbound] srcs = tuple(st.src[1] for st in c.src[0].src if st.op is Ops.STORE) # add the outputs to the call @@ -171,9 +172,25 @@ def replace_input_buffer(ctx:AllocCtx, b:UOp): ctx.replacements.append(b) return b.param_like(len(ctx.replacements)-1) -pm_replace_buf = PatternMatcher([ - # replace BUFFER with PARAM for cache key normalization (ALU addrspace buffers are Variables, they stay) - (UPat(Ops.BUFFER, src=(), name="b"), lambda ctx,b: replace_input_buffer(ctx, b) if b.addrspace is AddrSpace.GLOBAL else None), +# unbound BUFFERs get canonical scope-local id slots here so structurally identical calls hash identically for the +# schedule cache (fresh slots are all positive from the global counter; negative slots are already canonical) +def canonicalize_unbound_buffer(ctx:AllocCtx, b:UOp): + if b.arg.slot >= 0 and b not in ctx.unbound: ctx.unbound[b] = b.replace(arg=replace(b.arg, slot=-1-len(ctx.unbound))) + return ctx.unbound.get(b) + +def canonicalize_call_body(ctx:AllocCtx, c:UOp): + body = graph_rewrite(c.src[0], pm_canonicalize_unbound, ctx=ctx, bottom_up=True) + return c.replace(src=(body,)+c.src[1:]) if body is not c.src[0] else None + +pm_canonicalize_unbound = PatternMatcher([ + (UPat(Ops.CALL, name="c"), canonicalize_call_body), + (UPat(Ops.BUFFER, src=(), name="b"), lambda ctx,b: canonicalize_unbound_buffer(ctx, b) if b.is_unbound else None), +]) + +pm_replace_buf = pm_canonicalize_unbound+PatternMatcher([ + # replace BUFFER with PARAM for cache key normalization (ALU addrspace buffers are Variables, they stay, and unbound BUFFERs too) + (UPat(Ops.BUFFER, src=(), name="b"), lambda ctx,b: + replace_input_buffer(ctx, b) if b.addrspace is AddrSpace.GLOBAL and not b.is_unbound else None), # replace buffer views (SHRINK/BITCAST) with PARAM (only the views created by contiguous_mops_to_view) (UPat((Ops.SHRINK, Ops.BITCAST), name="b"), lambda ctx,b: replace_input_buffer(ctx, b) if b in ctx.views else None), # strip the stored value from bound Variables for cache key normalization, so different values hit same cache @@ -195,7 +212,7 @@ def transform_to_call(big_sink:UOp) -> tuple[UOp, dict[UOp, UOp]]: # final outputs of value calls materialize with fresh storage srcs:list[UOp] = [] for u in big_sink.src: - if u.op is Ops.AFTER and u.src[0].unsharded_base.op is Ops.RETURNED: + if u.op is Ops.AFTER and u.src[0].unsharded_base.is_unbound: # precompiled calls don't need this: transform_precompiled_call gives their outputs real buffers call = u.src[1] if not (call.op is Ops.CALL and call.arg is not None and call.arg.precompile and call.num_returned): @@ -209,7 +226,7 @@ def transform_to_call(big_sink:UOp) -> tuple[UOp, dict[UOp, UOp]]: # collect the stores (never entering call bodies) and map tagged AFTERs to their storage; tags are stripped at the end # copies to disk are stores to the disk buffer; bound Variables are call inputs and RETURNEDs are call outputs for u in big_sink.toposort(enter_calls=False): - if (u.op is Ops.COPY and on_disk(u)) or (u.op is Ops.AFTER and not u.is_bound_var and u.src[0].unsharded_base.op is not Ops.RETURNED): + if (u.op is Ops.COPY and on_disk(u)) or (u.op is Ops.AFTER and not u.is_bound_var and not u.src[0].unsharded_base.is_unbound): ctx.stores.append(u) if u.tag: ctx.buffer_map.update({t:graph_rewrite(u.src[0], pm_drop_after).shrink_to(t.shape) for t in u.tag}) ret = graph_rewrite(UOp.sink(*ctx.stores), pm_replace_buf+remove_all_tags, ctx=ctx, bottom_up=True, name="replace bufs").call(*ctx.replacements) diff --git a/tinygrad/uop/__init__.py b/tinygrad/uop/__init__.py index 060d5671ff..3fac89dbc3 100644 --- a/tinygrad/uop/__init__.py +++ b/tinygrad/uop/__init__.py @@ -37,9 +37,6 @@ class Ops(FastEnum): # vector creation / item selection STACK = auto() - # RETURNED is a placeholder for a buffer a call writes and returns: it's an input to the call and you AFTER on it - RETURNED = auto() - # hcq specific GETADDR = auto() diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 79cfb639f2..c53d55475c 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -168,7 +168,7 @@ def dtype_from_uop(op:Ops, src:tuple[UOp,...], arg:Any) -> DType: if not all(dtypes.is_int(x.dtype) or x.base.is_invalid for x in src): raise RuntimeError(f"shift operands must be int, got {[x.dtype for x in src]}") return src[0].dtype - case Ops.BUFFER | Ops.PARAM | Ops.RETURNED: + case Ops.BUFFER | Ops.PARAM: assert isinstance(arg, ParamArg), f"{op} must have ParamArg" return arg.dtype case Ops.BINARY: @@ -360,7 +360,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): case Ops.GETADDR: return () case Ops.RANGE | Ops.SPECIAL: return () case Ops.BINARY: return (len(self.arg),) - case Ops.BUFFER | Ops.PARAM | Ops.RETURNED: + case Ops.BUFFER | Ops.PARAM: # these don't have a shape input, they have a size in the arg: int gives shape (size,), None gives () if (img:=self.arg.image) is not None: return (img[0], img[1], 4) return () if self.arg.size is None else (self.arg.size,) @@ -526,8 +526,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass): num = next(ucount) # tags can contain UOps (callify tags nodes with their originals): store them as trace_nums, same as srcs tag = tuple(t.trace_num if isinstance(t, UOp) else t for t in self.tag) if isinstance(self.tag, tuple) else self.tag - # the trace must not retain the device Buffer: drop it from the stored arg (it would pin the memory and fail trace pickling) - arg = replace(self.arg, buffer=None) if isinstance(self.arg, ParamArg) and self.arg.buffer is not None else self.arg + # the trace must not retain the device Buffer: store a placeholder instead (the real one would pin memory and fail pickling), + # keeping bound and unbound buffers distinguishable in viz + arg = replace(self.arg, buffer=cast("Buffer", object())) if isinstance(self.arg, ParamArg) and self.arg.buffer is not None else self.arg uop_fields[num] = (self.op, tuple(s.trace_num for s in self.src), arg, tag)+((self.metadata,) if TRACEMETA>=2 else ()) return num @@ -536,22 +537,23 @@ class UOp(RandMixin, metaclass=UOpMetaClass): def sink(*srcs:UOp|None, **kwargs): # pylint: disable=no-self-argument return UOp(Ops.SINK, src=tuple([x for x in srcs if x is not None]), **kwargs) @staticmethod - def returned(slot:int, dtype:DType, shape:tuple[sint, ...]|sint|None=None, device=None, axis:int|None=None) -> UOp: - """create a RETURNED placeholder for a buffer a call writes and returns: it's an input to the call and you AFTER on it - like a normal buffer. its slot is its position among the call's srcs, which is its identity (identical slots merge) + def returned(dtype:DType, shape:tuple[sint, ...]|sint|None=None, device=None, axis:int|None=None) -> UOp: + """create an unbound BUFFER declaration for a buffer a call writes and returns: it's an input to the call and you AFTER on + it like a normal buffer. its identity is unique (minted from the global counter): outputs of different calls never alias like PARAM, the arg only stores the concrete max size: a shape is a view (RESHAPE/SHRINK/UNSHARD) on the flat placeholder""" if isinstance(shape, (int, UOp)): shape = (shape,) + slot = next(UOp.unique_num) # multi-device values have a per-shard sized storage wrapped in UNSHARD: the sharding lives in the graph, not the arg - if shape is None or len(shape) == 0: return UOp(Ops.RETURNED, arg=ParamArg(slot, dtype, None, device=device)) + if shape is None or len(shape) == 0: return UOp(Ops.BUFFER, arg=ParamArg(slot, dtype, None, device=device)) shp = tuple(s//len(device) if (i == axis and isinstance(device, tuple)) else s for i,s in enumerate(shape)) - ret = UOp(Ops.RETURNED, arg=ParamArg(slot, dtype, prod(to_max_shape(shp)), device=device)) + ret = UOp(Ops.BUFFER, arg=ParamArg(slot, dtype, prod(to_max_shape(shp)), device=device)) return ret.view_as(shp, axis) @property - def num_returned(self) -> int: return sum(x.unsharded_base.op is Ops.RETURNED for x in self.src[1:]) + def num_returned(self) -> int: return sum(x.unsharded_base.is_unbound for x in self.src[1:]) @property def returned_outputs(self) -> tuple[UOp, ...]: """the outputs of a value-producing call: an AFTER on each RETURNED input, usable like a normal buffer""" - return tuple(x.after(self) for x in self.src[1:] if x.unsharded_base.op is Ops.RETURNED) + return tuple(x.after(self) for x in self.src[1:] if x.unsharded_base.is_unbound) # legacy compatibility: TUPLE/GETTUPLE are gone. a tuple of values called is call_outputs, gettuple is returned_outputs[i] @staticmethod def maketuple(*srcs:UOp) -> _LegacyTupleValues: return _LegacyTupleValues(srcs) @@ -857,7 +859,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): return ret.after(ret.store(src.cast(ret.dtype))) @recursive_property def device(self) -> str|tuple[str, ...]|None: - if self.op in (Ops.PARAM, Ops.RETURNED): return self.arg.device + if self.op is Ops.PARAM: return self.arg.device if self.op is Ops.STAGE: return self.arg.device if self.op is Ops.AFTER: return self.src[0].device if self.op is Ops.MSELECT: @@ -877,7 +879,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): return self.device is None or self.dtype in dtypes.weaks @recursive_property def addrspace(self) -> AddrSpace|None: - if self.op in (Ops.PARAM, Ops.RETURNED): return self.arg.addrspace + if self.op is Ops.PARAM: return self.arg.addrspace if self.op is Ops.BUFFER: return self.arg.addrspace if self.op in {Ops.SPECIAL, Ops.RANGE, Ops.CONST}: return AddrSpace.ALU if self.op is Ops.LOAD: return AddrSpace.ALU # LOAD brings things into the ALU @@ -891,7 +893,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass): return None @property def buf_uop(self) -> UOp: - if self.op in {Ops.BUFFER, Ops.PARAM, Ops.RETURNED}: return self + if self.op in {Ops.BUFFER, Ops.PARAM}: return self if self.op is Ops.MSELECT: return self.src[0].buf_uop.mselect(self.arg) if self.op is Ops.MSTACK: return UOp(Ops.MSTACK, src=tuple(x.buf_uop for x in self.src)) if self.base.op is Ops.AFTER: return self.base.src[0].buf_uop.base @@ -922,7 +924,11 @@ class UOp(RandMixin, metaclass=UOpMetaClass): # TODO: this is confusing because UOp.variable('v', 0, 1, dtypes.weakfloat) is True for jit to work, but it doesn't have a buffer if self.op in {Ops.RESHAPE, Ops.UNSHARD, Ops.MSELECT}: return self.src[0].has_buffer_identity(after_ok) if after_ok and self.op == Ops.AFTER: return self.src[0].has_buffer_identity(after_ok) - return self.op in {Ops.BUFFER, Ops.PARAM} + return self.op in {Ops.BUFFER, Ops.PARAM} and not self.is_unbound + @property + def is_unbound(self) -> bool: + # an unbound GLOBAL BUFFER has no storage bound yet: it's a declaration of storage (call output, scheduler temp) + return self.op is Ops.BUFFER and isinstance(self.arg, ParamArg) and self.addrspace is AddrSpace.GLOBAL and self.arg.buffer is None def _base_buffer_is_realized(self) -> bool: """Walk through AFTER chain to find if the underlying buffer is realized (has allocated memory).""" @@ -1219,11 +1225,10 @@ class UOp(RandMixin, metaclass=UOpMetaClass): # the device defaults to the first device in the values or args, like srcs-based device resolution default_dev = next((x.device for x in itertools.chain(values, srcs) if x.device is not None), None) # the RETURNED storage has the resolved shape: substitute internal PARAMs in the shapes with corresponding args - def returned(o:UOp, i:int) -> UOp: - return UOp.returned(len(srcs)+i, o.dtype, None if (shp:=o._shape) is None else - tuple(graph_rewrite(s, _pm_resolve_params, srcs, walk=True) if isinstance(s, UOp) else s for s in shp), - dev if (dev:=o.device) is not None else default_dev, o.axis if isinstance(o.device, tuple) else None) - rets = tuple(returned(o, i) for i, o in enumerate(values)) + rets = tuple(UOp.returned(o.dtype, None if (shp:=o._shape) is None else + tuple(graph_rewrite(s, _pm_resolve_params, srcs, walk=True) if isinstance(s, UOp) else s for s in shp), + o.device if o.device is not None else default_dev, o.axis if isinstance(o.device, tuple) else None) + for o in values) # the body only knows PARAMs: the output PARAMs get the slots right after the input PARAM slots body = UOp.sink(*[v.param_like(len(srcs)+i).store(v) for i, v in enumerate(values)]) return UOp(Ops.CALL, src=(body,)+srcs+rets, arg=CallInfo(grad_fxn, name, precompile, precompile_backward, aux)) @@ -1689,8 +1694,9 @@ class RewriteContext: continue # no rewrite, process children then come back to rebuild stack.append((n, True)) - # calls with RETURNED inputs are always inlined into the enclosing graph, their bodies are never rewritten separately - if n.op is Ops.CALL and (n.num_returned or (not self.enter_calls and n.src[0].op in UOp._OPAQUE_CALL_BODIES)): + # program bodies (kernels, value calls) are never rewritten separately unless the rewrite explicitly enters + # calls; other call graphs (dtype-arg calls) are plain dataflow and always rewritten + if n.op is Ops.CALL and not self.enter_calls and n.src[0].op in UOp._OPAQUE_CALL_BODIES: self.replace[n.src[0]] = n.src[0] for x in reversed(n.src): if x not in self.replace: stack.append((x, False)) @@ -1728,10 +1734,9 @@ class RewriteContext: if n in waitlist: stack.extend(waitlist.pop(n)) continue stack.append((n, 1, new_n)) - # NOTE: CALLs are handled as a special case: the call body is not included in the graph_rewrite (a CALL of an - # address is not a body, its srcs are regular dataflow). calls with RETURNED inputs are always inlined into the - # enclosing graph, their bodies are never rewritten separately - if new_n.op is Ops.CALL and (new_n.num_returned or (not self.enter_calls and new_n.src[0].op in UOp._OPAQUE_CALL_BODIES)): + # NOTE: CALLs are handled as a special case: program bodies are not included in the graph_rewrite unless the + # rewrite explicitly enters calls (a CALL of an address is not a body, its srcs are regular dataflow) + if new_n.op is Ops.CALL and not self.enter_calls and new_n.src[0].op in UOp._OPAQUE_CALL_BODIES: self.replace[new_n.src[0]] = new_n.src[0] for x in reversed(new_n.src): if x in on_stack: continue @@ -1786,7 +1791,7 @@ def resolve_returned_after(r:UOp, t:UOp) -> UOp|None: """AFTER on a RETURNED placeholder extracts the call output value: the value of its matching store in a SINK body (called from patterns that bind t to a SINK)""" vals = [st.src[1] for st in t.src if st.op is Ops.STORE and st.src[0].unsharded_base is r.unsharded_base] \ - if r.unsharded_base.op is Ops.RETURNED else [] + if r.unsharded_base.is_unbound else [] return vals[0] if len(vals) == 1 else None remove_all_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)]) diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index 5b872cf093..d790699360 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -95,7 +95,7 @@ spec_shared = PatternMatcher([ # AFTER on Movement Op, PARAM, BUFFER, CONTIGUOUS, RETURNED, or another AFTER (UPat(Ops.AFTER, src=(UPat(GroupOp.Movement.union({Ops.PARAM, Ops.BUFFER, Ops.CONTIGUOUS, Ops.INDEX, - Ops.AFTER, Ops.UNSHARD, Ops.BITCAST, Ops.INS, Ops.RETURNED})),), + Ops.AFTER, Ops.UNSHARD, Ops.BITCAST, Ops.INS})),), allow_any_len=True), lambda: True), # CUSTOM (inline and non inline): the arg is the source string and the dtype it produces, void for a bare statement @@ -125,7 +125,7 @@ spec_shared = PatternMatcher([ # STORE: the target must be storage or a CONTIGUOUS realization point (or an AFTER/BITCAST/view of one); # CONTIGUOUS targets are written into the buffer the CONTIGUOUS creates. INDEX stores are checked above (UPat(Ops.STORE, dtypes.void, (UPat(name="x"), UPat())), lambda x: - True if (b:=x.storage_base).op in {Ops.BUFFER, Ops.PARAM, Ops.RETURNED, Ops.CONTIGUOUS} else None if b.op is Ops.INDEX else False), + True if (b:=x.storage_base).op in {Ops.BUFFER, Ops.PARAM, Ops.CONTIGUOUS} else None if b.op is Ops.INDEX else False), # WMMA has a (UPat(Ops.WMMA, src=(UPat(), UPat(), UPat()), name="x"), lambda x: isinstance(x.arg, tuple) and len(x.arg) == 5), @@ -140,7 +140,7 @@ spec_tensor = PatternMatcher([ # BUFFER (UPat(Ops.BUFFER, src=(), name="buf"), lambda buf: - (isinstance(buf.dtype, DType) and isinstance(buf.arg.size, int) and is_device(buf.arg.device)) + True if buf.is_unbound else (isinstance(buf.dtype, DType) and isinstance(buf.arg.size, int) and is_device(buf.arg.device)) if isinstance(buf.arg, ParamArg) and buf.addrspace is AddrSpace.GLOBAL else None), # a Variable is a 0-d ALU BUFFER with a value range and no device @@ -152,9 +152,6 @@ spec_tensor = PatternMatcher([ # CALL (UPat(Ops.CALL, dtypes.void, src=(UPat((Ops.SINK, Ops.LINEAR, Ops.PROGRAM, Ops.COPY, Ops.CUSTOM_FUNCTION)),), allow_any_len=True), lambda: True), - # RETURNED is a placeholder for a buffer a call writes and returns: it has a size in the arg, no shape input - (UPat(Ops.RETURNED, src=(), name="x"), lambda x: isinstance(x.arg, ParamArg)), - # SPECIAL is index before index lowering. custom_kernel currently has this (UPat(Ops.SPECIAL, src=(UPat(dtype=dtypes.weakint),), name="s"), lambda s: isinstance(s.arg, str)), diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index e4ab45ea77..e0e2107f88 100755 --- a/tinygrad/viz/serve.py +++ b/tinygrad/viz/serve.py @@ -51,7 +51,7 @@ uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0", Ops.WMMA: "#efefc0", Ops.UNSHARD: "#f6ccff", Ops.INS: "#eec4ff", **{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80", Ops.BUFFER: "#B0BDFF", Ops.GETADDR: "#9DB1F0", Ops.COPY: "#a040a0", Ops.CUSTOM_FUNCTION: "#bf71b6", - Ops.CALL: "#00B7C8", Ops.PARAM: "#14686F", Ops.RETURNED: "#C07788", Ops.SOURCE: "#c0c0c0", Ops.BINARY: "#404040", + Ops.CALL: "#00B7C8", Ops.PARAM: "#14686F", Ops.SOURCE: "#c0c0c0", Ops.BINARY: "#404040", Ops.LINEAR: "#7DF4FF", Ops.ALLREDUCE: "#ff40a0", Ops.MSELECT: "#d040a0", Ops.MSTACK: "#d040a0", Ops.CONTIGUOUS: "#FFC14D", Ops.STAGE: "#AC640D", Ops.REWRITE_ERROR: "#1a1b26", Ops.AFTER: "#8A7866", Ops.END: "#524C46"} @@ -163,9 +163,11 @@ def uop_to_json(data:VizData, x:UOp) -> dict[int, dict]: # limit SOURCE labels line count if u.op is Ops.SOURCE and len(lines:=label.split("\n")) > 40: label = "\n".join(lines[:30]) + "\n..." + if u.is_unbound: label += "\nUNBOUND" addrspace_color:str|None = None with soft_err(): addrspace_color = addrspace_colors.get(u.addrspace, None) if u.addrspace is not None else None - graph[id(u)] = {"label":label, "src":[(i,id(x)) for i,x in enumerate(u.src)], "exclude":u in excluded, "color":uops_colors.get(u.op, "#ffffff"), + color = "#C07788" if u.is_unbound else uops_colors.get(u.op, "#ffffff") + graph[id(u)] = {"label":label, "src":[(i,id(x)) for i,x in enumerate(u.src)], "exclude":u in excluded, "color":color, "ref":ref, "tag":repr(u.tag) if u.tag is not None else None, "addrspace":addrspace_color} return graph