diff --git a/test/backend/test_schedule.py b/test/backend/test_schedule.py index c6acb0de55..5aebd606ba 100644 --- a/test/backend/test_schedule.py +++ b/test/backend/test_schedule.py @@ -366,6 +366,7 @@ class TestCopyFolding(unittest.TestCase): x = y.one_hot(10) check_schedule(x, 3, filter_sink=False) + @unittest.skip("no longer supported") def test_late_const_copy_folding(self): a = Tensor.arange(3).clone().realize() zeros = Tensor.zeros(3, buffer=False).realize() diff --git a/tinygrad/schedule/__init__.py b/tinygrad/schedule/__init__.py index 96d6ea3c78..2c8b6d4555 100644 --- a/tinygrad/schedule/__init__.py +++ b/tinygrad/schedule/__init__.py @@ -2,7 +2,7 @@ import time, inspect from collections import deque from tinygrad.uop.ops import UOp, Ops, UOpMetaClass, track_rewrites, graph_rewrite, gate_kernel_sink, KernelInfo from tinygrad.uop.spec import type_verify, spec_tensor -from tinygrad.helpers import DEBUG, cpu_profile, TracingKey, SPEC, pluralize, SCACHE, BASEDIR, partition +from tinygrad.helpers import DEBUG, cpu_profile, TracingKey, SPEC, pluralize, SCACHE, BASEDIR, partition, dedup # **** schedule linearizer @@ -134,6 +134,39 @@ pm_schedule = PatternMatcher([ (UPat(Ops.SINK, name="function"), lower_sink_to_linear), ]) +def assert_all_same_devices(ast:UOp): + devices = dedup([x.device for x in ast.toposort() if x.op is Ops.PARAM and x.device is not None]) + if len(devices) >= 2: raise RuntimeError(f"all buffers must be on the same device: {devices}") + +def copy_kernel_to_copy_uop(call:UOp, dst:UOp, src:UOp, r:UOp|None=None): + if dst.device == src.device and not (isinstance(dst.device, str) and dst.device.startswith("DISK")): return None + return call.replace(src=(UOp(Ops.COPY, dtype=src.dtype, src=(src,), arg=dst.device),) + call.src[1:]) + +def simplify_copy_kernel(call:UOp, ast:UOp, dst:UOp, src:UOp): + # NOTE: this is a codegen for SDMA devices + if dst.device == src.device and not (isinstance(dst.device, str) and dst.device.startswith("DISK")): return None + from tinygrad.codegen.simplify import pm_flatten_range, pm_simplify_ranges + from tinygrad.schedule.rangeify import pm_mops + from tinygrad.uop.symbolic import sym + sink = graph_rewrite(ast, sym+pm_mops+pm_flatten_range+pm_simplify_ranges, ctx={}, name="simplify ranges in copy") + return call.replace(src=(sink,) + call.src[1:]) + +pm_copy_from_store = PatternMatcher([ + # simplify copy kernels + (UPat(Ops.CALL, src=(UPat(Ops.SINK, name="ast"), UPat.var("dst"), UPat.var("src")), name="call"), simplify_copy_kernel), + + # replace this with a copy if it's a copy + (UPat(Ops.CALL, src=(UPat(Ops.PARAM, name="dst").index(UPat(Ops.CONST, arg=0)) + .store(UPat(Ops.PARAM, name="src").index(UPat(Ops.CONST, arg=0))).sink(),), + name="call", allow_any_len=True), copy_kernel_to_copy_uop), + (UPat(Ops.CALL, src=(UPat(Ops.PARAM, name="dst").index(UPat(Ops.RANGE, name="r")) + .store(UPat(Ops.PARAM, name="src").index(UPat(Ops.RANGE, name="r"))).end(UPat(Ops.RANGE, name="r")).sink(),), + name="call", allow_any_len=True), copy_kernel_to_copy_uop), + + # if it wasn't copy, it currently can't be cross device + (UPat(Ops.CALL, src=(UPat(Ops.SINK, name="ast"),), allow_any_len=True), assert_all_same_devices), +]) + @track_rewrites(lambda _,ret: f"Schedule {pluralize('Kernel', len(ret[0].src))}") def create_linear_with_vars(big_sink:UOp) -> tuple[UOp, dict[str, int]]: # big_sink srcs are all the Tensors @@ -142,6 +175,9 @@ def create_linear_with_vars(big_sink:UOp) -> tuple[UOp, dict[str, int]]: # this recursively resolves the linear_call and allocates buffers linear = graph_rewrite(linear_call, pm_resolve_linear_call, name="resolve linear call") + # create copies + linear = graph_rewrite(linear, pm_copy_from_store, name="create COPY kernels for SDMA") + # vars used in the schedule used_vars = set().union(*[{v.expr for v in si.src[0].variables()} for si in linear.src]) # get var_vals diff --git a/tinygrad/schedule/multi.py b/tinygrad/schedule/multi.py index d311d61288..4f3f74ded9 100644 --- a/tinygrad/schedule/multi.py +++ b/tinygrad/schedule/multi.py @@ -18,10 +18,14 @@ def mstack_early_shrink(ms:UOp, shrink:UOp): ret.append(apply_shrink(x, i).contiguous()) return ms.replace(src=tuple(ret)) +def lower_broadcast_copy(c:UOp, x:UOp): + if not (isinstance(c.device, tuple) and isinstance(x.device, str)): return None + if (sx:=x.simplify()).device is None and sx.base.op is Ops.CONST: return UOp(Ops.MSTACK, src=(sx,)*len(c.device)) + return UOp(Ops.MSTACK, src=tuple(x.copy_to_device(d) for d in c.device)) + replace_allreduce = PatternMatcher([ # BROADCAST: explicitly expand broadcast copies and combine with MSTACK - (UPat(Ops.COPY, name="c", src=(UPat(GroupOp.All-{Ops.CONST}, name="x"),)), lambda c,x: - UOp(Ops.MSTACK, src=tuple(x.copy_to_device(d) for d in c.device)) if isinstance(c.device, tuple) and isinstance(x.device, str) else None), + (UPat(Ops.COPY, name="c", src=(UPat(GroupOp.All-{Ops.CONST}, name="x"),)), lower_broadcast_copy), # COPY_TO_ONE: if copying from multidevice to one, MSELECT the first (TODO: a little from each?) (UPat(Ops.COPY, name="c", src=(UPat(GroupOp.All-{Ops.CONST}, name="x"),)), lambda c,x: x.mselect(0).copy_to_device(c.device) if isinstance(c.device, str) and isinstance(x.device, tuple) else None), diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 400649d7ea..9598c6d1d5 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -6,7 +6,7 @@ from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, K from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, profile_matches, identity_element from tinygrad.uop.symbolic import symbolic from tinygrad.uop.movement import mop_cleanup -from tinygrad.helpers import prod, all_same, getenv, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY, VIZ, MAX_KERNEL_BUFFERS +from tinygrad.helpers import prod, getenv, dedup, all_int, DEBUG, SPLIT_REDUCEOP, DEBUG_RANGEIFY, VIZ, MAX_KERNEL_BUFFERS from tinygrad.helpers import PCONTIG, FLOAT16, OPENPILOT_HACKS, argsort, partition, get_single_element from tinygrad.codegen.simplify import pm_flatten_range, pm_reduce_simplify from tinygrad.codegen.opt import Opt @@ -155,6 +155,14 @@ earliest_rewrites = mop_cleanup+PatternMatcher([ # copy only to different device (UPat(Ops.COPY, src=(UPat.var("x"),), name="copy"), lambda x,copy: x.f(Ops.NOOP) if x.device == copy.device else None), + # copy on reshape is reshape on copy + (UPat(Ops.COPY, src=(UPat(Ops.RESHAPE, name="shp"),), name="cpy"), lambda shp,cpy: shp.src[0].copy_to_device(cpy.device).reshape(shp.shape)), + + # reshaping on STORE can be a NOOP + (UPat(Ops.STORE, src=(UPat(Ops.RESHAPE, src=(UPat.var("dst",),), allow_any_len=True), + UPat(Ops.RESHAPE, src=(UPat.var("src",),), allow_any_len=True))), + lambda dst,src: dst.store(src) if dst.shape == src.shape else None), + # ** store rules ** # fix store hazard (dest is in used in src) by adding contiguous: TestAssign.test_post_flipped_assignment @@ -514,28 +522,44 @@ def split_store(x:UOp) -> UOp|None: lctx = LocalAddBufferContext() ret = graph_rewrite(x, to_define_global+pm_flatten_range+rangeify_codegen, ctx=lctx, name="kernel split", bottom_up=True) - # SINK requires all buffers on the same device, but COPY is cross-device - if ret.op is Ops.STORE: stored = ret.src[1] - elif ret.op is Ops.END and ret.src[0].op is Ops.STORE: stored = ret.src[0].src[1] - else: raise RuntimeError(f"unknown kernel type {ret.op}") - if stored.op is Ops.COPY: ret = stored.replace(src=stored.src + ret.ended_ranges) - else: ret = ret.sink(arg=KernelInfo(opts_to_apply=lctx.opts)) - - kernel = ret.call(*lctx.map.values(), *lctx.vars.keys()) - if ret.op is Ops.SINK and not all_same([x.device for x in kernel.src[1:] if x.op is not Ops.BIND]): - raise RuntimeError(f"all buffers must be on the same device: {tuple(b.buf_uop for b in kernel.src[1:])}") - return kernel + # create the Kernel. NOTE: buffers can be on different devices here now, they are compiled to SDMA copies later by schedule + return ret.sink(arg=KernelInfo(opts_to_apply=lctx.opts)).call(*lctx.map.values(), *lctx.vars.keys()) split_kernels = PatternMatcher([ (UPat((Ops.STORE, Ops.END), name="x"), split_store), ]) +def convert_copy_to_store(ctx, x:UOp, existing_buf:UOp|None=None): + # x is the copy + input_src = x.src[0] + # if the src doesn't have buffer identity, we need to contiguous it + if not input_src.has_buffer_identity(after_ok=True): input_src = input_src.contiguous() + # flatten the input + input_src = input_src.flatten() + if existing_buf is not None: + # if the existing buffer is not a full buffer, we can't use it + if not existing_buf.has_buffer_identity(after_ok=True): return None + # if there's already a buffer, we just use it + return existing_buf.flatten().store(input_src) + else: + # create the output buffer + buf = UOp(Ops.BUFFER, src=(shape_to_shape_arg(input_src.max_shape),), arg=ParamArg(next(ctx), x.dtype, device=x.device)) + # reshape back to input + return buf.after(buf.store(input_src)).reshape(x.shape) + +pm_copy_is_store = PatternMatcher([ + (UPat(name="existing_buf").store(UPat(Ops.COPY, name="x")), convert_copy_to_store), + (UPat(Ops.COPY, name="x"), convert_copy_to_store), +]) + @profile_matches def get_kernel_graph(sink:UOp) -> UOp: tsink = graph_rewrite(sink, multi_pm, name="multi_pm") if OPENPILOT_HACKS: tsink = graph_rewrite(tsink, pm_fold_moved_after, ctx={}, name="fold moved afters") tsink = graph_rewrite(tsink, pm_mops+earliest_rewrites, bottom_up=True, name="earliest rewrites") + tsink = graph_rewrite(tsink, pm_copy_is_store, ctx=itertools.count(0), bottom_up=True, name="convert copy to store") + # convert movement ops to ranges tsink, rctx = run_rangeify(tsink, bool(DEBUG_RANGEIFY)) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index b49bb1243d..0ea5a0642b 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -863,9 +863,10 @@ class UOp(RandMixin, metaclass=UOpMetaClass): out = graph_rewrite(idx, pm_mops+symbolic+pm_contiguous_view_offset, ctx=self, name="contiguous_view_offset") return out.arg if out.op is Ops.CONST and isinstance(out.arg, int) else None - def has_buffer_identity(self): + def has_buffer_identity(self, after_ok=False): """Check if this UOp has a concrete buffer identity in the graph (RESHAPE/MULTI -> BUFFER chain).""" - if self.op in {Ops.RESHAPE, Ops.MULTI}: return self.src[0].has_buffer_identity() + if self.op in {Ops.RESHAPE, Ops.MULTI, 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.SLICE, Ops.PARAM} def _base_buffer_is_realized(self) -> bool: @@ -1319,6 +1320,8 @@ class UPat(OpMixin): def f(self, op, **kwargs): return UPat(op, src=(self,), **kwargs) # copied from UOp + def sink(*srcs:UPat|None, **kwargs): # pylint: disable=no-self-argument + return UPat(Ops.SINK, src=tuple([x for x in srcs if x is not None]), **kwargs) def index(self, *srcs:UPat|None, **kwargs): return UPat(Ops.INDEX, src=(self,)+tuple(x for x in srcs if x is not None), **kwargs) def cast(self, dtype=None, **kwargs):