From cddc4dcfc06de79519b451c1991c580cf24db706 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Tue, 11 Aug 2026 10:58:28 -0700 Subject: [PATCH] split kernels --- tinygrad/schedule/rangeify2.py | 28 +++++++++++++++++++++------- 1 file changed, 21 insertions(+), 7 deletions(-) diff --git a/tinygrad/schedule/rangeify2.py b/tinygrad/schedule/rangeify2.py index e07b03153d..ee85f62298 100644 --- a/tinygrad/schedule/rangeify2.py +++ b/tinygrad/schedule/rangeify2.py @@ -138,20 +138,34 @@ pm_range_migration = PatternMatcher([ # *** split into kernels *** -def _split_graph(ctx, u:UOp): - ctx.append(u) - return u.param_like(len(ctx)-1) -pm_split_graph = PatternMatcher([(UPat((Ops.PARAM, Ops.AFTER), name="u"), _split_graph),]) +@dataclass +class SplitCtx: + call_args:list = field(default_factory=list) + range_number:int = -1 + +def _split_graph(ctx:SplitCtx, u:UOp) -> UOp: + ctx.call_args.append(u) + return u.param_like(len(ctx.call_args)-1) + +def _renumber_range(ctx:SplitCtx, u:UOp) -> UOp: + ctx.range_number += 1 + return u.replace(arg=(ctx.range_number, u.arg[-1])) + +pm_split_graph = PatternMatcher([ + (UPat((Ops.PARAM, Ops.AFTER), name="u"), _split_graph), + (UPat(Ops.RANGE, name="u"), _renumber_range), +]) def split_store(x:UOp) -> UOp: - kernel_splits = [] - ret = graph_rewrite(x, pm_split_graph, ctx=kernel_splits, name="split kernel", bottom_up=True, walk=True) - return ret.sink(arg=KernelInfo()).call(*kernel_splits) + ret = graph_rewrite(x, pm_split_graph, ctx:=SplitCtx(), name="split kernel", bottom_up=True, walk=True) + return ret.sink(arg=KernelInfo()).call(*ctx.call_args) split_kernels = PatternMatcher([ (UPat((Ops.STORE, Ops.END), name="x"), split_store), ]) +# *** main rangeify *** + @rewrite_group(new_ctx=False) def get_kernel_graph(sink:UOp) -> UOp: # TODO: multi should just be part of rangeify