From aa98aff4cd58aa8e914cd887826fd75310ae9c57 Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Fri, 18 Apr 2025 08:59:18 +0100 Subject: [PATCH] don't use ops name, just keep sink (#9922) * don't use ops name, just keep sink * fix test * endif sink --- test/external/external_benchmark_schedule.py | 4 ++-- test/test_linearizer.py | 3 ++- test/test_uop_graph.py | 6 +++++- tinygrad/codegen/kernel.py | 2 +- tinygrad/codegen/linearize.py | 8 +++----- tinygrad/ops.py | 5 +++-- tinygrad/renderer/cstyle.py | 4 ++-- tinygrad/renderer/llvmir.py | 4 ++-- tinygrad/renderer/ptx.py | 4 ++-- tinygrad/runtime/ops_python.py | 4 ++-- tinygrad/spec.py | 2 +- tinygrad/upat.py | 2 +- tinygrad/viz/serve.py | 2 +- 13 files changed, 27 insertions(+), 23 deletions(-) diff --git a/test/external/external_benchmark_schedule.py b/test/external/external_benchmark_schedule.py index 8398c48ba6..1f1c9923b2 100644 --- a/test/external/external_benchmark_schedule.py +++ b/test/external/external_benchmark_schedule.py @@ -15,7 +15,7 @@ if __name__ == "__main__": for p in nn.state.get_parameters(mdl): p.replace(Tensor.empty(p.shape)) img = Tensor.empty(64, 3, 224, 224) - PROFILE = getenv("PROFILE", 0) + PROFILE = getenv("PYPROFILE", 0) FORWARD_ONLY = getenv("FORWARD_ONLY", 0) SCHEDULE_ONLY = getenv("SCHEDULE_ONLY", 0) @@ -50,7 +50,7 @@ if __name__ == "__main__": rewritten_uops.append(full_graph_rewrite(u, k.opts)) uops = rewritten_uops if getenv("LINEARIZE", 1): - with Profiling(PROFILE >= 2): + with Profiling(PROFILE >= 2, frac=0.5): with Timing("***** model linearize in "): uops = [linearize_uop(u) for u in uops] print(sum(len(u) for u in uops)) if getenv("SRC", 0): diff --git a/test/test_linearizer.py b/test/test_linearizer.py index a74a748497..03b6a9804c 100644 --- a/test/test_linearizer.py +++ b/test/test_linearizer.py @@ -846,7 +846,8 @@ class TestLinearizer(unittest.TestCase): sink = UOp(Ops.SINK, src=(store,)) load_t = Tensor.full(load.st_arg.shape, 1).contiguous().realize() k = helper_linearizer_ast(sink, [load_t], wanna_output=[load_t.numpy().sum()])[1] - self.assertEqual(k.uops[-1].op, Ops.ENDIF) + self.assertEqual(k.uops[-2].op, Ops.ENDIF) + self.assertEqual(k.uops[-1].op, Ops.SINK) self.assertLess(k.uops.index([x for x in k.uops if x.op is Ops.STORE][-1]), k.uops.index(k.uops[-1])) def test_two_nested_range(self): diff --git a/test/test_uop_graph.py b/test/test_uop_graph.py index 6207c94359..9cc5b0393e 100644 --- a/test/test_uop_graph.py +++ b/test/test_uop_graph.py @@ -17,7 +17,11 @@ simple_pm = PatternMatcher([ ((UPat.var('x') + UPat.cvar('c1')) + UPat.cvar('c2'), lambda x,c1,c2: x + (c1.arg+c2.arg)), ]) -def to_uops_list(u:List[UOp]) -> List[UOp]: return linearize_uop(full_graph_rewrite(UOp.sink(*u))) +def to_uops_list(u:List[UOp]) -> List[UOp]: + # we strip the SINK here for legacy reasons + ret = linearize_uop(full_graph_rewrite(UOp.sink(*u))) + assert ret[-1].op is Ops.SINK + return ret[:-1] class TestGraphRewriteEfficiency(unittest.TestCase): def test_create_many_uops(self): diff --git a/tinygrad/codegen/kernel.py b/tinygrad/codegen/kernel.py index 680d70cbed..76d70aa4eb 100644 --- a/tinygrad/codegen/kernel.py +++ b/tinygrad/codegen/kernel.py @@ -565,7 +565,7 @@ class Kernel: def to_program(self, name_override:Optional[str]=None, ast_transform:Optional[Callable]=None) -> ProgramSpec: self.linearize(name_override, ast_transform) - assert self.uops[0].op is Ops.NAME, "first uop must be name" + assert self.uops[-1].op is Ops.SINK, "last uop must be sink" src = self.opts.render(self.uops) if CAPTURE_PROCESS_REPLAY: diff --git a/tinygrad/codegen/linearize.py b/tinygrad/codegen/linearize.py index 86001c1149..17dee0be7c 100644 --- a/tinygrad/codegen/linearize.py +++ b/tinygrad/codegen/linearize.py @@ -6,7 +6,7 @@ from tinygrad.spec import type_verify from tinygrad.dtype import dtypes, PtrDType from tinygrad.helpers import dedup, flatten, partition -DONT_PLACE_IN_BLOCK = {Ops.NAME, Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_VAR, Ops.SPECIAL, Ops.CONST, *GroupOp.Block} +DONT_PLACE_IN_BLOCK = {Ops.DEFINE_GLOBAL, Ops.DEFINE_LOCAL, Ops.DEFINE_VAR, Ops.SPECIAL, Ops.CONST, *GroupOp.Block} def disp(y:UOp) -> str: if y.op is Ops.BLOCKSTART: return "w"+disp(y.src[0]) @@ -74,9 +74,7 @@ def append_to_block(ctx:tuple[dict[UOp, tuple[UOp, ...]], dict[UOp, list[UOp]]], return UOp(Ops.BLOCK, dtypes.void, tuple(dedup(list(old_blocks.values())+new_srcs)), BasicBlock(x.arg.ctx, tuple(to_append)+x.arg.lst)) make_basic_blocks = PatternMatcher([ - (UPat(Ops.SINK, name="x"), - lambda x: UOp(Ops.BLOCK, src=x.src+((UOp(Ops.NAME, arg=x.arg.name),) if x.arg is not None else ()), arg=BasicBlock((), (x,)))), - (UPat(Ops.BLOCK, name="x"), append_to_block), + (UPat(Ops.SINK, name="x"), lambda x: UOp(Ops.BLOCK, src=x.src, arg=BasicBlock((), (x,)))), (UPat(Ops.BLOCK, name="x"), append_to_block), ]) def block_merge(ctx, x:UOp): @@ -132,7 +130,7 @@ def block_finalize(block:UOp): _uops += block.arg.lst # strip the SINK assert _uops[-1].op is Ops.SINK, "doesn't end with SINK" - return UOp(Ops.BLOCK, arg=BasicBlock((), tuple(_uops[:-1]))) + return UOp(Ops.BLOCK, arg=BasicBlock((), tuple(_uops))) pm_block_finalize = PatternMatcher([(UPat(Ops.BLOCK, name="block"), block_finalize)]) diff --git a/tinygrad/ops.py b/tinygrad/ops.py index 72aeadbdd0..2dd45321ee 100644 --- a/tinygrad/ops.py +++ b/tinygrad/ops.py @@ -93,7 +93,7 @@ class MathTrait(SimpleMathTrait): # the order of these Ops controls the order of the toposort class Ops(FastEnum): # uops that aren't rendered - NAME = auto(); SINK = auto(); CONTIGUOUS = auto(); CONTIGUOUS_BACKWARD = auto(); DETACH = auto(); KERNEL = auto(); UNIQUE = auto() # noqa: E702 + SINK = auto(); CONTIGUOUS = auto(); CONTIGUOUS_BACKWARD = auto(); DETACH = auto(); KERNEL = auto(); UNIQUE = auto() # noqa: E702 # MetaOps COPY = auto(); BUFFER_VIEW = auto() # noqa: E702 @@ -285,7 +285,8 @@ class UOp(MathTrait, metaclass=UOpMetaClass): def get_children_map(self) -> dict[UOp, dict[UOp, None]]: ret: dict[UOp, dict[UOp, None]] = {} for u in self.toposort: - for s in u.src: ret.setdefault(s, {})[u] = None + ret[u] = {} + for s in u.src: ret[s][u] = None return ret @functools.cached_property diff --git a/tinygrad/renderer/cstyle.py b/tinygrad/renderer/cstyle.py index b9ca112931..f77ba84ec9 100644 --- a/tinygrad/renderer/cstyle.py +++ b/tinygrad/renderer/cstyle.py @@ -132,8 +132,8 @@ class CStyleLanguage(Renderer): c: defaultdict[str, int] = defaultdict(int) name = "test" for u in uops: - if u.op is Ops.NAME: - name = u.arg + if u.op is Ops.SINK: + if u.arg is not None: name = u.arg.name continue if u.op in (Ops.DEFINE_GLOBAL, Ops.DEFINE_VAR): r[u] = f"data{u.arg}" if u.op is Ops.DEFINE_GLOBAL else u.arg[0] diff --git a/tinygrad/renderer/llvmir.py b/tinygrad/renderer/llvmir.py index 1946b73eb9..e292be317e 100644 --- a/tinygrad/renderer/llvmir.py +++ b/tinygrad/renderer/llvmir.py @@ -164,8 +164,8 @@ class LLVMRenderer(Renderer): name = "test" for u in uops: - if u.op is Ops.NAME: - name = u.arg + if u.op is Ops.SINK: + if u.arg is not None: name = u.arg.name continue if u.op in (Ops.DEFINE_GLOBAL, Ops.DEFINE_VAR): r[u] = f"%data{u.arg}" if u.op is Ops.DEFINE_GLOBAL else f"%{u.arg[0]}" diff --git a/tinygrad/renderer/ptx.py b/tinygrad/renderer/ptx.py index f07fa0d4cb..bdfb44bd12 100644 --- a/tinygrad/renderer/ptx.py +++ b/tinygrad/renderer/ptx.py @@ -176,8 +176,8 @@ class PTXRenderer(Renderer): name = "test" for u in uops: - if u.op is Ops.NAME: - name = u.arg + if u.op is Ops.SINK: + if u.arg is not None: name = u.arg.name continue if u.op is Ops.VECTORIZE: r[u] = [cast(str,r[x]) for x in u.src] diff --git a/tinygrad/runtime/ops_python.py b/tinygrad/runtime/ops_python.py index 3d66448c0c..6d127b8eff 100644 --- a/tinygrad/runtime/ops_python.py +++ b/tinygrad/runtime/ops_python.py @@ -40,7 +40,7 @@ class PythonProgram: loop_ends: dict[int, int] = {} while i < len(self.uops): uop, dtype, idp, arg = self.uops[i] - void_ops = {Ops.STORE, Ops.ENDRANGE, Ops.BARRIER, Ops.IF, Ops.ENDIF, Ops.NAME} + void_ops = {Ops.STORE, Ops.ENDRANGE, Ops.BARRIER, Ops.IF, Ops.ENDIF, Ops.SINK} if uop is Ops.DEFINE_ACC: idp = [idp[0]] inp = [ul[v] for v in idp if self.uops[v][0] not in void_ops] dtp = [dl[v] for v in idp if self.uops[v][0] not in void_ops] @@ -60,7 +60,7 @@ class PythonProgram: loop_ends[idp[0]] = i i = idp[0] continue - if uop in (Ops.BARRIER, Ops.IF, Ops.ENDIF, Ops.NAME): + if uop in (Ops.BARRIER, Ops.IF, Ops.ENDIF, Ops.SINK): # in the python emulator, the warp is always in sync i += 1 continue diff --git a/tinygrad/spec.py b/tinygrad/spec.py index ed180a2620..6273642ca4 100644 --- a/tinygrad/spec.py +++ b/tinygrad/spec.py @@ -132,7 +132,7 @@ spec = PatternMatcher([ # NOTE: for testing, we let sinks be anything #(UPat(Ops.SINK, src=UPat(Ops.STORE)), lambda: True), - (UPat((Ops.NAME, Ops.SINK), dtypes.void), lambda: True), + (UPat(Ops.SINK, dtypes.void), lambda: True), (UPat((Ops.NOOP, Ops.CUSTOMI, Ops.CUSTOM)), lambda: True), # PTX LOAD/STORE diff --git a/tinygrad/upat.py b/tinygrad/upat.py index 0edabe0a0c..0ab3dfe7e6 100644 --- a/tinygrad/upat.py +++ b/tinygrad/upat.py @@ -139,7 +139,7 @@ def _final_render(x:UOp, has_ctx:bool, depth=1) -> list[str]: def _get_code(self:UPat, has_ctx:bool): ret = _get_clause(self, UOp(Ops.NOOP, arg="uop")) try: - ret = graph_rewrite(ret, pm_proc) + ret = graph_rewrite(ret, pm_proc, name="process UPat") dyn_lookup: dict[str, Any] = {} out = graph_rewrite(ret, pm_renderer, ctx=dyn_lookup, name="compile UPat") rendered = _final_render(out, has_ctx) diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index 77dc38a7b1..0e330f35fd 100755 --- a/tinygrad/viz/serve.py +++ b/tinygrad/viz/serve.py @@ -14,7 +14,7 @@ uops_colors = {Ops.LOAD: "#ffc0c0", Ops.STORE: "#87CEEB", Ops.CONST: "#e0e0e0", Ops.RANGE: "#c8a0e0", Ops.ASSIGN: "#e0ffc0", Ops.BARRIER: "#ff8080", Ops.IF: "#c8b0c0", Ops.SPECIAL: "#c0c0ff", Ops.INDEX: "#e8ffa0", Ops.WMMA: "#efefc0", Ops.VIEW: "#C8F9D4", Ops.MULTI: "#f6ccff", Ops.KERNEL: "#3e7f55", Ops.IGNORE: "#00C000", **{x:"#D8F9E4" for x in GroupOp.Movement}, **{x:"#ffffc0" for x in GroupOp.ALU}, Ops.THREEFRY:"#ffff80", Ops.BUFFER_VIEW: "#E5EAFF", - Ops.BLOCK: "#C4A484", Ops.BLOCKEND: "#C4A4A4", Ops.BUFFER: "#B0BDFF", Ops.COPY: "#a040a0", Ops.NAME:"#808080"} + Ops.BLOCK: "#C4A484", Ops.BLOCKEND: "#C4A4A4", Ops.BUFFER: "#B0BDFF", Ops.COPY: "#a040a0"} # VIZ API