diff --git a/tinygrad/schedule/rangeify2.py b/tinygrad/schedule/rangeify2.py index ee85f62298..e7096c79ac 100644 --- a/tinygrad/schedule/rangeify2.py +++ b/tinygrad/schedule/rangeify2.py @@ -118,6 +118,11 @@ def index_on_stack(stack:UOp, idx:UOp): for k in range(len(srcs)-2, -1, -1): ret = r0.eq(k).where(srcs[k], ret) return ret +def walk_mop(u:UOp): + if u.op in GroupOp.Movement or u.op is Ops.INDEX: return u.src[0] + assert u.op == Ops.AFTER + return u + pm_range_migration = PatternMatcher([ # INDEX without src is nothing (UPat(Ops.INDEX, src=(UPat.var('x'),)), lambda x: x), @@ -131,9 +136,14 @@ pm_range_migration = PatternMatcher([ # handle movement ops on INDEX (UPat(GroupOp.Movement, name="r").index(name="idx", allow_any_len=True), _mop_index), (UPat(Ops.STACK, name="stack").index(name="idx", allow_any_len=True), index_on_stack), + # move movement ops and INDEX after AFTER + (UPat(GroupOp.Movement|{Ops.INDEX}, name="r").after(name="a", allow_any_len=True), + lambda r,a: UOp(r.op, src=(a.replace(src=(r.src[0],)+a.src[1:]),)+r.src[1:], arg=r.arg)), # pass index through elementwise (UPat(GroupOp.Elementwise, name="b").index(name="idx", allow_any_len=True), lambda b,idx: b.replace(src=tuple(s.index(*idx.src[1:]) for s in b.src))), + # remove movement ops from SINK. TODO: should be generic + (UPat(Ops.SINK, name="s"), lambda s: s.replace(src=tuple(walk_mop(u) for u in s.src))), ]) # *** split into kernels *** @@ -144,6 +154,7 @@ class SplitCtx: range_number:int = -1 def _split_graph(ctx:SplitCtx, u:UOp) -> UOp: + assert len(u.shape) == 1, "rangeify needs to reduce to a single idx" ctx.call_args.append(u) return u.param_like(len(ctx.call_args)-1) @@ -166,6 +177,10 @@ split_kernels = PatternMatcher([ # *** main rangeify *** +debug_tag_factor = PatternMatcher([ + (UPat(GroupOp.All, name="x"), lambda ctx,x: x.rtag(ctx[x]) if x.tag is None else None), +]) + @rewrite_group(new_ctx=False) def get_kernel_graph(sink:UOp) -> UOp: # TODO: multi should just be part of rangeify @@ -177,6 +192,17 @@ def get_kernel_graph(sink:UOp) -> UOp: # TODO: add safe STAGEs to never duplicate compute + # we compute the number of times a buffer is consumed, everything starts with 0 + consumes = {tsink:0} + for u in reversed(tsink.toposort()): + assert u in consumes + if u.op is Ops.STORE: consumes[u] = 1 + if u.op is Ops.EXPAND: consumes[u] *= u.max_numel() // u.src[0].max_numel() + for s in u.src[1:] if u.op is Ops.STORE else u.src: + if s not in consumes: consumes[s] = 0 + consumes[s] += consumes[u] + if VIZ: graph_rewrite(tsink, debug_tag_factor, ctx=consumes, name="view consumes tags", bottom_up=True) + # simple rangeify tsink = graph_rewrite(tsink, pm_range_creation+pm_range_migration, ctx=itertools.count(0), bottom_up=True, name="simple rangeify")