diff --git a/tinygrad/engine/schedule.py b/tinygrad/engine/schedule.py index 054bc8922e..b9d22f6779 100644 --- a/tinygrad/engine/schedule.py +++ b/tinygrad/engine/schedule.py @@ -154,6 +154,7 @@ def contig_to_assign(ctx:dict[UOp,UOp|None], x:UOp): # 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), ]) schedule_cache: dict[bytes, tuple[list[ExecItem], UOp]] = {} @track_rewrites(lambda _,ret: f"Schedule {pluralize('Kernel', len(ret[1]))}") @@ -162,8 +163,8 @@ def complete_create_schedule_with_vars(big_sink:UOp) -> tuple[dict[UOp, UOp], li 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), ]) + # we assign for any contiguous or anything on the sink that's not a BUFFER, a CONST, or already ASSIGN + buffer_map: dict[UOp, UOp|None] = {x.base:None for x in big_sink.src if x.base.op not in {Ops.CONST, Ops.BUFFER, Ops.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())