From c4d565b591da85f911eee84ba69426a3cc7bf9f0 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Fri, 15 Aug 2025 18:46:04 -0700 Subject: [PATCH] improve names --- tinygrad/schedule/rangeify.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/tinygrad/schedule/rangeify.py b/tinygrad/schedule/rangeify.py index 149255db75..e82750ef3b 100644 --- a/tinygrad/schedule/rangeify.py +++ b/tinygrad/schedule/rangeify.py @@ -2,7 +2,7 @@ from typing import Any from dataclasses import dataclass, field from tinygrad.dtype import dtypes, AddrSpace, PtrDType from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, RewriteNotReady, _substitute -from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv +from tinygrad.helpers import argsort, prod, all_same, pluralize, getenv, colored from tinygrad.uop.symbolic import symbolic_simple from tinygrad.schedule.kernelize import Kernel @@ -361,10 +361,13 @@ to_define_global = PatternMatcher([ def split_store(x:UOp): if len(x.ranges): return None - shape = tuple([r.vmax+1 for r in x.src[2:]]) - name = "k_"+'_'.join([str(s) for s in shape]) + store_rngs = x.src[2:] + ctx = LocalAddBufferContext() ret = graph_rewrite(x, to_define_global, ctx=ctx, name="kernel split", bottom_up=True) + rng = [u for u in ret.toposort() if u.op is Ops.RANGE] + name = "k_"+colored('_', 'BLACK').join([colored(str(s.vmax+1), "blue") if s in store_rngs else colored(str(s.vmax+1), "red") for s in rng]) + ret = ret.sink(arg=KernelInfo(name=name)) if ret.op is Ops.STORE else ret kernel = UOp(Ops.KERNEL, src=tuple([x[0] for x in ctx.map.values()])+tuple(ctx.vars.keys()), arg=Kernel(ret, ())) return kernel.src[0].assign(kernel)