better

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())