From a5d3b54f47dfebdc50bfc7f1d4648a4dfba0880a Mon Sep 17 00:00:00 2001 From: George Hotz Date: Thu, 14 Aug 2025 13:52:14 -0700 Subject: [PATCH] work --- tinygrad/schedule/kernelize.py | 16 ++++++++++++---- tinygrad/schedule/rangeify.py | 9 ++++----- 2 files changed, 16 insertions(+), 9 deletions(-) diff --git a/tinygrad/schedule/kernelize.py b/tinygrad/schedule/kernelize.py index 295a98deec..72e3608b35 100644 --- a/tinygrad/schedule/kernelize.py +++ b/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") diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 1862e9039a..3f1957b1af 100644 --- a/tinygrad/schedule/rangeify.py +++ b/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