work

tinygrad/schedule/kernelize.py

@@ -330,7 +330,8 @@ new_fixups = mops_merge+PatternMatcher([
 # *** store splitting
 
 def split_load(ctx:list[UOp], s:UOp):
-  if len(s.src) == 1 or s.src[0].src[0].op is not Ops.DEFINE_GLOBAL: return None
+  if len(s.src) == 1: return None
+  ctx.extend(s.src[1:])
   return s.replace(src=s.src[0:1])
 
 def debuf(ctx:list[int], b:UOp):
@@ -338,18 +339,25 @@ def debuf(ctx:list[int], b:UOp):
   ctx.append(b)
   return ret
 
+def maybe_copy(s:UOp):
+  if s.src[1].op is not Ops.COPY: return None
+  return UOp(Ops.COPY)
+  #b0 = s.src[0].src[0]
+  #b1 = s.src[1].src[0].src[0].src[0]
+  #return UOp(Ops.COPY, src=(b0, b1))
+
 to_define_global = PatternMatcher([
   (UPat(Ops.BUFFER, name="b"), debuf),
   (UPat(Ops.LOAD, name="s"), split_load),
+  (UPat(Ops.STORE, name="s"), maybe_copy),
 ])
 
 def split_store(x:UOp):
   shape = tuple([r.vmax+1 for r in x.src[2:]])
   name = "k_"+'_'.join([str(s) for s in shape])
-
   b = x.src[0].src[0]
   ctx = []
-  ret = graph_rewrite(x, to_define_global, ctx=ctx, name="* kernel split")
+  ret = graph_rewrite(x, to_define_global, ctx=ctx, name="* kernel split", bottom_up=True)
   ret = ret.sink(arg=KernelInfo(name=name))
   kernel = UOp(Ops.KERNEL, src=(b,)+tuple(ctx), arg=Kernel(ret, ()))
   return b.assign(kernel)
@@ -380,7 +388,7 @@ def get_kernelize_map(sink:UOp) -> dict[UOp, UOp]:
     tensor_map = graph_rewrite_map(tensor_map[sink], pm_children, ctx=ChildrenContext(), bottom_up=True, input_map=tensor_map, name="* children")
     tensor_map = graph_rewrite_map(tensor_map[sink], pm_rangeify, ctx=RangeifyContext(), bottom_up=True, input_map=tensor_map, name="* rangeify")
     tensor_map = graph_rewrite_map(tensor_map[sink], pm_add_buffers, ctx=AddBufferContext(), bottom_up=True, input_map=tensor_map, name="* buffer")
-    tensor_map = graph_rewrite_map(tensor_map[sink], split_kernels, bottom_up=True, input_map=tensor_map, name="* split kernels")
+    tensor_map = graph_rewrite_map(tensor_map[sink], split_kernels, input_map=tensor_map, name="* split kernels")
 
     if getenv("VIZ"): graph_rewrite(tensor_map[sink], PatternMatcher([]), name="View Kernel Graph")
     #rsink = graph_rewrite(rsink, sym, name="* symbolic")

tinygrad/schedule/rangeify.py

@@ -7,6 +7,8 @@ from tinygrad.helpers import argsort, prod, all_same
 rangeify_fixups = PatternMatcher([
   # all contiguous on SINK
   (UPat(Ops.SINK, name="x"), lambda x: x.replace(src=tuple([s.contiguous() if s.op not in {Ops.CONTIGUOUS, Ops.CONST} else s for s in x.src]))),
+  # all contiguous on COPY
+  (UPat(Ops.COPY, name="x"), lambda x: x.replace(tag=1).contiguous() if x.tag is None else None),
   # double contiguous merge
   (UPat(Ops.CONTIGUOUS, name="c2", src=(UPat(Ops.CONTIGUOUS, name="c1"))), lambda c1,c2: c1 if c1.arg is None and c2.arg is None else None),
   # const
@@ -201,11 +203,10 @@ pm_rangeify = pm_mops+PatternMatcher([
         allow_any_len=True, name="idx"), indexed_endrange),
 
   # move MAP through elementwise ALU / reduce. these are the items with cost
-  (UPat(Ops.INDEX, src=(UPat(GroupOp.Elementwise.union({Ops.STORE, Ops.ASSIGN})),), allow_any_len=True, name="x"),
+  (UPat(Ops.INDEX, src=(UPat(GroupOp.Elementwise.union({Ops.STORE, Ops.ASSIGN, Ops.COPY, Ops.DEVICE})),), allow_any_len=True, name="x"),
    lambda x: x.src[0].replace(src=tuple([s.index(*x.src[1:]) for s in x.src[0].src]))),
   (UPat(Ops.INDEX, src=(UPat(Ops.REDUCE_AXIS, name="red"),), allow_any_len=True, name="idx"), map_reduce),
 
-
   # CONTIGUOUS on ASSIGN is STORE
   # TODO: tag in UPat?
   (UPat(Ops.CONTIGUOUS, src=(UPat(Ops.ASSIGN, name="a"),), name="c", allow_any_len=True),
@@ -237,7 +238,7 @@ def add_load(ctx:AddBufferContext, x:UOp, b:UOp, idx:UOp):
 def add_load_on_store(ctx:AddBufferContext, x:UOp, st:UOp):
   rngs = x.src[1:]
   shape = tuple([r.vmax+1 for r in rngs])
-  return st.src[0].src[0].shrink(((0,prod(shape)),)).reshape(shape).index(*rngs).load(st)
+  return st.src[0].src[0].shrink(((0,prod(shape)),)).reshape(shape).index(*rngs, dtype=x.dtype.ptr(size=st.src[0].src[0].size)).load(st)
 
 pm_add_buffers = pm_mops+PatternMatcher([
   (UPat(Ops.CONTIGUOUS, name="x"), add_store),
@@ -245,8 +246,6 @@ pm_add_buffers = pm_mops+PatternMatcher([
   (UPat(Ops.INDEX, src=(UPat(Ops.BUFFER, name="b"), UPat(name="idx")), name="x"), add_load),
   (UPat(Ops.INDEX, src=(UPat(Ops.STORE, name="st"),), allow_any_len=True, name="x"), add_load_on_store),
   (UPat(Ops.BIND, name="b"), lambda b: b.src[0]),
-  # HACK: ignore copy
-  (UPat(Ops.COPY, name="x"), lambda x: x.src[0]),
   # CONST can't have axes. remove srcs when we idx
   (UPat(Ops.INDEX, src=(UPat(Ops.CONST, name="c"),)), lambda c: c.replace(src=())),
   # HACK: consts shouldn't have srcs by here