From 1884f67a3934816671c1ea2a2f50b7864a93090e Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Mon, 4 May 2026 11:44:13 -0700 Subject: [PATCH] simplify full_rewrite_to_sink spec (#16035) * simplify full_rewrite_to_sink spec * test cleanups --- test/helpers.py | 9 ++++++--- test/null/test_const_folding.py | 6 +++--- test/null/test_graph_rewrite.py | 10 +++++----- test/null/test_simplify_valid_idx.py | 12 ++++++------ tinygrad/codegen/__init__.py | 18 ++++++++---------- 5 files changed, 28 insertions(+), 27 deletions(-) diff --git a/test/helpers.py b/test/helpers.py index 80227b2ae1..6b7c1add0e 100644 --- a/test/helpers.py +++ b/test/helpers.py @@ -17,11 +17,14 @@ from tinygrad.codegen.late.linearizer import linearize slow = unittest.skipUnless(os.getenv("RUN_SLOW"), "slow test, set RUN_SLOW=1 to run") from tinygrad.runtime.ops_python import PythonProgram, PythonRenderer, PythonCompiler -def get_uops(sink:UOp, ren:Renderer|None=None) -> list[UOp]: - """Extract linearized UOps from a sink. Test helper that only does linearization (no render).""" +def full_rewrite(sink:UOp, ren:Renderer|None=None) -> UOp: if ren is None: ren = Renderer(Target()) if sink.arg is None: sink = sink.replace(arg=KernelInfo()) - full_sink = full_rewrite_to_sink(sink, ren, optimize=sink.tag is None) + return full_rewrite_to_sink(sink, ren, optimize=sink.tag is None) + +def get_uops(sink:UOp, ren:Renderer|None=None) -> list[UOp]: + """Extract linearized UOps from a sink. Test helper that only does linearization (no render).""" + full_sink = full_rewrite(sink, ren) return line_rewrite(linearize(full_sink), pm_linearize_cleanups) def replace_opts(ast:UOp, opts:list) -> UOp: return ast.replace(arg=replace(ast.arg, opts_to_apply=tuple(opts))) diff --git a/test/null/test_const_folding.py b/test/null/test_const_folding.py index 1d244609af..aedb3eb6e1 100644 --- a/test/null/test_const_folding.py +++ b/test/null/test_const_folding.py @@ -2,7 +2,7 @@ import unittest, itertools, math from tinygrad import Tensor, dtypes, Context from tinygrad.dtype import DType, ConstType from tinygrad.uop.ops import Ops, UOp -from tinygrad.codegen import full_rewrite_to_sink +from test.helpers import full_rewrite import numpy as np def _check_ast_count(desired_count:int, t:Tensor): @@ -103,7 +103,7 @@ class TestBitcastConstFolding(unittest.TestCase): def t(cases: dict[DType, ConstType]): for (from_dt, from_v), (to_dt, to_v) in itertools.product(cases.items(), cases.items()): if not math.isnan(from_v): - r = full_rewrite_to_sink(UOp.const(from_dt, from_v).bitcast(to_dt).sink()).src[0] + r = full_rewrite(UOp.const(from_dt, from_v).bitcast(to_dt).sink()).src[0] self.assertEqual(r.op, Ops.CONST, msg:=f"{from_dt} -> {to_dt} ({from_v} -> {to_v})") self.assertEqual(r.dtype, to_dt, msg) np.testing.assert_equal(r.arg, to_v, msg) @@ -127,7 +127,7 @@ class TestBitcastConstFolding(unittest.TestCase): def test_vec_bitcast(self): with Context(SPEC=0): - r = full_rewrite_to_sink(UOp.const(dtypes.int32.vec(3), (-1, -2**31, 75)).bitcast(dtypes.uint32.vec(3)).sink()).src[0] + r = full_rewrite(UOp.const(dtypes.int32.vec(3), (-1, -2**31, 75)).bitcast(dtypes.uint32.vec(3)).sink()).src[0] self.assertEqual(r.op, Ops.STACK) self.assertEqual(r.dtype, dtypes.uint32.vec(3)) self.assertEqual(tuple(x.arg for x in r.src), (2**32-1, 2**31, 75)) diff --git a/test/null/test_graph_rewrite.py b/test/null/test_graph_rewrite.py index fa99f200c0..8af3132491 100644 --- a/test/null/test_graph_rewrite.py +++ b/test/null/test_graph_rewrite.py @@ -2,13 +2,13 @@ import unittest, math from tinygrad import dtypes from tinygrad.helpers import all_same, Context from tinygrad.uop.ops import GroupOp, UOp, Ops, exec_alu, PatternMatcher, TrackedPatternMatcher, UPat -from tinygrad.codegen import full_rewrite_to_sink +from test.helpers import full_rewrite from hypothesis import given, strategies as strat # Helper function to apply the graph rewrite @Context(SPEC=0) def apply_rewrite(expr): - return full_rewrite_to_sink(expr.sink()).src[0] + return full_rewrite(expr.sink()).src[0] def evaluate_uop(uop, variables): if uop.op == Ops.CONST: @@ -151,7 +151,7 @@ class TestModuloAndDivisionFolding(unittest.TestCase): class TestEdgeCasesAndSpecialOperations(unittest.TestCase): def test_full_graph_rewrite_transcendental_edge_cases(self): - optimized_sink = full_rewrite_to_sink(UOp.const(dtypes.float32, -1.0).log2().sink(UOp.const(dtypes.float32, 0.0).reciprocal())) + optimized_sink = full_rewrite(UOp.const(dtypes.float32, -1.0).log2().sink(UOp.const(dtypes.float32, 0.0).reciprocal())) optimized_log2_neg, optimized_recip_zero = optimized_sink.src self.assertTrue(math.isnan(optimized_log2_neg.arg), f"Expected NaN for log2(-1.0), got {optimized_log2_neg.arg}") self.assertTrue(math.isinf(optimized_recip_zero.arg) and optimized_recip_zero.arg > 0, @@ -160,14 +160,14 @@ class TestEdgeCasesAndSpecialOperations(unittest.TestCase): @unittest.skip("broken") def test_full_graph_rewrite_modulo_negative_dividend(self): x_var_uop = UOp.variable('x', -5, -1) - optimized_sink = full_rewrite_to_sink((x_var_uop % 3).sink()) + optimized_sink = full_rewrite((x_var_uop % 3).sink()) for x_value in range(-5, 0): self.assertEqual(x_value % 3, evaluate_uop(optimized_sink.src[0], {'x': x_value})) @unittest.skip("broken") def test_full_graph_rewrite_division_negative_divisor(self): x_var_uop = UOp.variable('x', 1, 5) - optimized_sink = full_rewrite_to_sink((x_var_uop // -2).sink()) + optimized_sink = full_rewrite((x_var_uop // -2).sink()) for x_value in range(1, 6): self.assertEqual(x_value // -2, evaluate_uop(optimized_sink.src[0], {'x': x_value})) diff --git a/test/null/test_simplify_valid_idx.py b/test/null/test_simplify_valid_idx.py index 50de6109b9..2174ce5d7c 100644 --- a/test/null/test_simplify_valid_idx.py +++ b/test/null/test_simplify_valid_idx.py @@ -1,10 +1,10 @@ import unittest, itertools -from tinygrad.codegen import full_rewrite_to_sink from tinygrad.dtype import dtypes from tinygrad.uop.ops import UOp, Ops from tinygrad.uop.symbolic import simplify_valid from tinygrad.helpers import Context +from test.helpers import full_rewrite from test.null.test_uop_symbolic import check_uop_against_string def get_gated_load_uop(valid:UOp, idx:UOp): @@ -48,7 +48,7 @@ class TestHelpers(unittest.TestCase): class TestValidIdxSimplification(unittest.TestCase): def check(self, load, sidx, svalid, extra=()): with Context(NOOPT=1, SPEC=0): - load = full_rewrite_to_sink(UOp.sink(load, *extra)).src[0] + load = full_rewrite(UOp.sink(load, *extra)).src[0] idx, valid = load.src[0].src[1], load.src[0].src[2] check_uop_against_string(self, idx, sidx) check_uop_against_string(self, valid, svalid) @@ -217,7 +217,7 @@ class TestValidIdxSimplification(unittest.TestCase): class TestImageSimplification(unittest.TestCase): def check(self, load, svalid, sidx0, sidx1): with Context(NOOPT=1, SPEC=0): - load = full_rewrite_to_sink(load.sink()).src[0] + load = full_rewrite(load.sink()).src[0] idx = load.src[0].src[1] self.assertEqual(idx.op, Ops.STACK) self.assertEqual(len(idx.src), 2) @@ -287,7 +287,7 @@ class TestImageSimplification(unittest.TestCase): # empty -> invalid load = get_load_image_uop(shape, (gidx0<8) & (gidx0<8).ne(True), idx) with Context(NOOPT=1, SPEC=0): - load = full_rewrite_to_sink(load.sink()).src[0] + load = full_rewrite(load.sink()).src[0] self.assertEqual(load.op, Ops.STACK) self.assertEqual(load.dtype.count, 4) @@ -508,7 +508,7 @@ class TestUnfoldableImage(unittest.TestCase): with Context(SPEC=0): lidx = Special("lidx", 2) load = UOp(Ops.LOAD, dtypes.float, (UOp(Ops.PARAM, dtypes.imagef((10, 10, 4)), arg=0).index(lidx, ptr=True), UOp.const(dtypes.float, 0))) - res = full_rewrite_to_sink(load.sink()).src[0] + res = full_rewrite(load.sink()).src[0] self.assertEqual(res.src[0].src[0].dtype, dtypes.float.ptr(400)) class TestDropTrueGate(unittest.TestCase): @@ -528,7 +528,7 @@ class TestDropTrueGate(unittest.TestCase): class TestRangeShrink(unittest.TestCase): def get_ranges(self, sink): with Context(NOOPT=1, SPEC=0): - result = full_rewrite_to_sink(sink) + result = full_rewrite(sink) return [u for u in result.toposort() if u.op is Ops.RANGE] def test_range_shrink_single_guard(self): diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index ade53701a5..d078c09c0b 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -2,7 +2,7 @@ from typing import cast from dataclasses import replace import itertools from tinygrad.helpers import DISABLE_FAST_IDIV, DEVECTORIZE, TRANSCENDENTAL, SPEC, DEBUG, VIZ, IMAGE, NOOPT, EMULATED_DTYPES, NOLOCALS, USE_TC -from tinygrad.helpers import ALLOW_TF32, TracingKey, Context, Target, panic +from tinygrad.helpers import ALLOW_TF32, TracingKey, Context, panic from tinygrad.uop.ops import PatternMatcher, graph_rewrite, UOp, pm_lower_index_dtype, Ops, UPat, track_rewrites, KernelInfo, ProgramInfo from tinygrad.uop.render import pyrender from tinygrad.uop.spec import type_verify, program_spec, kernel_spec @@ -21,15 +21,13 @@ from tinygrad.codegen.simplify import pm_simplify_ranges, pm_flatten_range, pm_s from tinygrad.schedule.rangeify import pm_add_buffers_local, rangeify_codegen, pm_mops, pm_syntactic_sugar, pm_store_ranges from tinygrad.codegen.late.linearizer import CFGContext, pm_split_ends, pm_add_control_flow, linearize -def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True, beam:int=0) -> UOp: - if ren is None: ren = Renderer(Target()) - - if VIZ: graph_rewrite(sink, PatternMatcher([]), name="View Base AST") - if DEBUG >= 5: print(pyrender(sink)) - if SPEC: type_verify(sink, kernel_spec) +def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: + if VIZ: graph_rewrite(ast, PatternMatcher([]), name="View Base AST") + if DEBUG >= 5: print(pyrender(ast)) + if SPEC: type_verify(ast, kernel_spec) # preprocess - sink = graph_rewrite(sink, pm_mops+pm_syntactic_sugar+pm_store_ranges, ctx=itertools.count(1000), name="early movement ops", bottom_up=True) + sink = graph_rewrite(ast, pm_mops+pm_syntactic_sugar+pm_store_ranges, ctx=itertools.count(1000), name="early movement ops", bottom_up=True) # first we optimize if optimize: @@ -46,7 +44,7 @@ def full_rewrite_to_sink(sink:UOp, ren:Renderer|None=None, optimize:bool=True, b sink = graph_rewrite(sink, pm_flatten_range+pm_simplify_ranges, ctx={}, name="simplify ranges") # do postrange optimization, BEAM or hand_coded_optimizations - sink = apply_opts(sink, ren, beam=beam) + sink = apply_opts(sink, ren, beam=ast.arg.beam) # ** expander (expand_rewrite) ** sink = graph_rewrite(sink, sym+pm_move_where_on_load, name="postopt symbolic") @@ -171,7 +169,7 @@ def do_to_program(ast:UOp, renderer:Renderer) -> UOp: if ast.op is Ops.PROGRAM: prg = ast elif ast.op is Ops.SINK: assert isinstance(ast.arg, KernelInfo), "requires KernelInfo on arg to to_program" - full_sink = full_rewrite_to_sink(ast, renderer, optimize=ast.tag is None, beam=ast.arg.beam) + full_sink = full_rewrite_to_sink(ast, renderer, optimize=ast.tag is None) prg = UOp(Ops.PROGRAM, src=(full_sink, UOp(Ops.DEVICE, arg=renderer.target.device)), arg=ProgramInfo.from_sink(full_sink)) else: raise RuntimeError(f"can't call to_program on {ast.op}") if not isinstance(prg.arg, ProgramInfo): prg = prg.replace(arg=ProgramInfo.from_sink(prg.src[0]))