diff --git a/tinygrad/engine/schedule.py b/tinygrad/engine/schedule.py index 8b7a7d4204..054bc8922e 100644 --- a/tinygrad/engine/schedule.py +++ b/tinygrad/engine/schedule.py @@ -2,6 +2,7 @@ import time from typing import cast from collections import deque from tinygrad.uop.ops import UOp, Ops, buffers, UOpMetaClass, track_rewrites, PatternMatcher, UPat, graph_rewrite, graph_rewrite_map, gate_kernel_sink +from tinygrad.uop.ops import _remove_all_tags, GroupOp from tinygrad.uop.spec import type_verify, tensor_spec from tinygrad.device import Buffer, MultiBuffer from tinygrad.helpers import DEBUG, cpu_profile, TracingKey, SPEC, flatten, pluralize, SCACHE, Metadata @@ -144,12 +145,29 @@ pm_post_sched_cache = PatternMatcher([ (UPat(Ops.BIND, src=(UPat(Ops.DEFINE_VAR),), name="b"), lambda ctx,b: ctx.get(b)), ]) +# rewrite all contiguous to assign +def contig_to_assign(ctx:dict[UOp,UOp|None], x:UOp): + # for contiguous or in buffer_map explicitly + if not (x.op is Ops.CONTIGUOUS or (x in ctx and ctx[x] is None)): return None + # not for symbolic shape + if any([not isinstance(s, int) for s in x.shape]): return None + # not sure why the ctx isn't enough, but tag fixes it + ctx[x] = buffer = UOp.new_buffer(x.device, x.size, x.dtype).reshape(x.shape) + return buffer.assign(x.src[0] if x.op is Ops.CONTIGUOUS else x.rtag()) + schedule_cache: dict[bytes, tuple[list[ExecItem], UOp]] = {} @track_rewrites(lambda _,ret: f"Schedule {pluralize('Kernel', len(ret[1]))}") def complete_create_schedule_with_vars(big_sink:UOp) -> tuple[dict[UOp, UOp], list[ExecItem], dict[str, int]]: # big_sink srcs are all the Tensors st = time.perf_counter() + # new preschedule stuff + buffer_map: dict[UOp, UOp|None] = {x.base:None for x in big_sink.src if x.base.op is not Ops.CONST} + pm_contig_to_assign = PatternMatcher([ (UPat(GroupOp.All, name="x"), contig_to_assign), ]) + big_sink = graph_rewrite(big_sink, pm_contig_to_assign, ctx=buffer_map, bottom_up=True, name="contig to assign") + big_sink = graph_rewrite(big_sink, _remove_all_tags) + assert all(x is not None for x in buffer_map.values()) + # replace BUFFERs with PARAMs, CONSTs UNIQUE with LUNIQUE, strip BIND values for cache key, extract var_vals input_buffers: dict[UOp, UOp] = {} var_vals: dict[str, int] = {} @@ -219,4 +237,6 @@ def complete_create_schedule_with_vars(big_sink:UOp) -> tuple[dict[UOp, UOp], li f" | {len(UOpMetaClass.ucache)} uops in cache") used_vars = set().union(*[{v.arg[0] for v in si.ast.variables()} for si in schedule]) - return tensor_map, schedule, {k:v for k,v in var_vals.items() if k in used_vars} + #return tensor_map, schedule, {k:v for k,v in var_vals.items() if k in used_vars} + # tensor_map isn't used anymore + return cast(dict[UOp, UOp], buffer_map), schedule, {k:v for k,v in var_vals.items() if k in used_vars} diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 36846ef617..f52228a0cd 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -257,28 +257,9 @@ class Tensor(OpMixin): """ big_sink = UOp.sink(*[x.uop for x in (self,)+lst]) - # rewrite all contiguous to assign - buffer_map: dict[UOp, UOp|None] = {x.base:None for x in big_sink.src} - from tinygrad.uop.ops import graph_rewrite, PatternMatcher, UPat, GroupOp, _remove_all_tags - def contig_to_assign(ctx:dict[UOp,UOp|None], x:UOp): - # for contiguous or in buffer_map explicitly - if not (x.op is Ops.CONTIGUOUS or (x in ctx and ctx[x] is None)): return None - # not for symbolic shape - if any([not isinstance(s, int) for s in x.shape]): return None - # not sure why the ctx isn't enough, but tag fixes it - ctx[x] = buffer = UOp.new_buffer(x.device, x.size, x.dtype).reshape(x.shape) - return buffer.assign(x.src[0] if x.op is Ops.CONTIGUOUS else x.rtag()) - - pm_contig_to_assign = PatternMatcher([ (UPat(GroupOp.All, name="x"), contig_to_assign), ]) - big_sink = graph_rewrite(big_sink, pm_contig_to_assign, ctx=buffer_map, bottom_up=True, name="contig to assign") - big_sink = graph_rewrite(big_sink, _remove_all_tags) - - assert all(x is not None for x in buffer_map.values()) - _apply_map_to_tensors(cast(dict[UOp, UOp], buffer_map), name="Apply Preallocated Buffers") - # this is where the schedule cache should go becomes_map, schedule, var_vals = complete_create_schedule_with_vars(big_sink) - #_apply_map_to_tensors(becomes_map, name="Apply Schedule Map") + _apply_map_to_tensors(becomes_map, name="Apply Schedule Map") return schedule, var_vals def schedule(self, *lst:Tensor) -> list[ExecItem]: