From 701a632907adee683e98fd7831654a0df0e5daec Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Mon, 27 Oct 2025 17:37:13 +0800 Subject: [PATCH] move VECTORIZE/CONST (#12942) --- test/test_uops.py | 10 ++++++++-- tinygrad/codegen/opt/search.py | 2 +- tinygrad/engine/realize.py | 2 +- tinygrad/uop/__init__.py | 1 + tinygrad/uop/ops.py | 4 ++-- tinygrad/uop/symbolic.py | 5 +++-- tinygrad/viz/serve.py | 2 +- 7 files changed, 17 insertions(+), 9 deletions(-) diff --git a/test/test_uops.py b/test/test_uops.py index 4b3102b78f..c55ae0ff27 100644 --- a/test/test_uops.py +++ b/test/test_uops.py @@ -548,10 +548,16 @@ class TestUopsObject(unittest.TestCase): class TestUOpRender(unittest.TestCase): def test_render_vectorize_same(self): u = UOp(Ops.VECTORIZE, dtype=dtypes.int.vec(3), src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0))) - self.assertEqual(u.render(), "{0, ...}") + self.assertEqual(u.render(simplify=False), "{0, ...}") def test_render_vectorize_different(self): u = UOp(Ops.VECTORIZE, dtype=dtypes.int.vec(3), src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1), UOp.const(dtypes.int, 2))) - self.assertEqual(u.render(), "{0,1,2}") + self.assertEqual(u.render(simplify=False), "{0,1,2}") + def test_render_vectorize_same_simplified(self): + u = UOp(Ops.VECTORIZE, dtype=dtypes.int.vec(3), src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 0))) + self.assertEqual(u.render(), "0") + def test_render_vectorize_different_simplified(self): + u = UOp(Ops.VECTORIZE, dtype=dtypes.int.vec(3), src=(UOp.const(dtypes.int, 0), UOp.const(dtypes.int, 1), UOp.const(dtypes.int, 2))) + self.assertEqual(u.render(), "(0, 1, 2)") if __name__ == '__main__': unittest.main(verbosity=2) diff --git a/tinygrad/codegen/opt/search.py b/tinygrad/codegen/opt/search.py index 8c0ff422d5..afce7048ee 100644 --- a/tinygrad/codegen/opt/search.py +++ b/tinygrad/codegen/opt/search.py @@ -137,7 +137,7 @@ def beam_search(lin:Scheduler, rawbufs:list[Buffer], amt:int, allow_test_size=Tr min_progress = getenv("BEAM_MIN_PROGRESS", 0.01)/1e6 if BEAM_DEBUG: print("BEAM_SEARCH:") - print('\n'.join(pyrender(lin.ast.replace(arg=None)))) + print(pyrender(lin.ast.replace(arg=None))) if DEBUG >= 2: print(f" 0.00s: from 1 -> 1 actions {lin.colored_shape()}") try: diff --git a/tinygrad/engine/realize.py b/tinygrad/engine/realize.py index db46840d95..b7286ac7ad 100644 --- a/tinygrad/engine/realize.py +++ b/tinygrad/engine/realize.py @@ -26,7 +26,7 @@ def get_program(ast:UOp, renderer:Renderer|None=None, opts:list[Opt]|None=None) """ if getenv("VIZ"): graph_rewrite(ast, PatternMatcher([]), name="View Base AST") - if DEBUG >= 5: print('\n'.join(pyrender(ast))) + if DEBUG >= 5: print(pyrender(ast)) # linearize if renderer is None: renderer = Device.default.renderer diff --git a/tinygrad/uop/__init__.py b/tinygrad/uop/__init__.py index 55d0f9e996..cad734d6d0 100644 --- a/tinygrad/uop/__init__.py +++ b/tinygrad/uop/__init__.py @@ -3,6 +3,7 @@ from enum import auto, IntEnum, Enum # wrapper around IntEnum that preserves Enum.__str__ and makes auto() unique across all FastEnum subclasses class FastEnum(IntEnum): def __str__(self): return Enum.__str__(self) + def __repr__(x): return str(x) @staticmethod def _generate_next_value_(_, __, ___, last_values): return 1 + max([0, *last_values, *[max(c) for c in FastEnum.__subclasses__()]]) diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index a404b7ae7a..734155ccd3 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -1256,7 +1256,7 @@ pm_pyrender = PatternMatcher([ ]) @Context(SPEC=0) -def pyrender(ast:UOp) -> list[str]: +def pyrender(ast:UOp) -> str: cmap = ast.get_consumer_map() to_render = set() for u in ast.toposort(): @@ -1271,7 +1271,7 @@ def pyrender(ast:UOp) -> list[str]: if u not in to_render: continue ret.append(f"c{len(ret)} = {u.substitute(rep).render(simplify=False, pm=pm_pyrender+renderer)}") rep[u] = UOp(Ops.NOOP, arg=f"c{len(ret)-1}") - return ret[0:-1] + ["ast ="+ret[-1].split("=", 1)[1]] + return "\n".join(ret[0:-1] + ["ast ="+ret[-1].split("=", 1)[1]]) # *** what was symbolic.py *** diff --git a/tinygrad/uop/symbolic.py b/tinygrad/uop/symbolic.py index 0852a12cb1..a677876548 100644 --- a/tinygrad/uop/symbolic.py +++ b/tinygrad/uop/symbolic.py @@ -382,6 +382,8 @@ symbolic = symbolic_simple+commutative+PatternMatcher([ tuple(flatten([(y,) if y.op in {Ops.RANGE, Ops.IF, Ops.STORE, Ops.KERNEL, Ops.BARRIER, Ops.END, Ops.UNROLL} else y.src for y in x.src[1:]])))), # after with 1 src is just src[0] (UPat(Ops.AFTER, src=(UPat.var("s"),)), lambda s: s), + # VECTORIZE/CONST + (UPat(Ops.VECTORIZE, src=UPat(Ops.CONST), name="vec"), lambda vec: UOp.const(vec.dtype, tuple(x.arg for x in vec.src))), ])+gep_pushing symbolic_flat = symbolic+PatternMatcher([ @@ -511,8 +513,7 @@ sym = symbolic_flat+pm_simplify_valid+PatternMatcher([ # LOAD/STORE -> NOOP (UPat.var('x').store(UPat.var('x').load(), allow_any_len=True), lambda x: None if x.dtype.addrspace != AddrSpace.REG else x.src[0].src[0]), (UPat(Ops.LOAD, src=(UPat.cvar('c'))), lambda c: c), - # VECTORIZE/CONST, VECTORIZE/GEP - (UPat(Ops.VECTORIZE, src=UPat(Ops.CONST), name="vec"), lambda vec: UOp.const(vec.dtype, tuple(x.arg for x in vec.src))), + # VECTORIZE/GEP (UPat(Ops.VECTORIZE, src=UPat(Ops.GEP, src=(UPat.var("x"),)), name="vec"), lambda vec,x: x.gep(tuple(y.arg[0] for y in vec.src))), # reorder ALU/VECTORIZE (UPat(GroupOp.ALU, src=(UPat(Ops.VECTORIZE, src=UPat(name='x')), UPat(Ops.VECTORIZE, src=UPat(name='y'))), name='alu'), diff --git a/tinygrad/viz/serve.py b/tinygrad/viz/serve.py index dc4b1498d9..9e4abcc2f0 100755 --- a/tinygrad/viz/serve.py +++ b/tinygrad/viz/serve.py @@ -54,7 +54,7 @@ def shape_to_str(s:tuple[sint, ...]): return "(" + ','.join(srender(x) for x in def mask_to_str(s:tuple[tuple[sint, sint], ...]): return "(" + ','.join(shape_to_str(x) for x in s) + ")" def pystr(u:UOp, i:int) -> str: if isinstance(trace.keys[i].ret, ProgramSpec): - try: return "\n".join(pyrender(u)) + try: return pyrender(u) except Exception: pass return str(u)