From a7c693d2fd13650ae15a7793586b666c276fc5ea Mon Sep 17 00:00:00 2001 From: chenyu Date: Wed, 22 Jul 2026 15:53:11 -0400 Subject: [PATCH] rework pm_lower_weakint [pr] (#17136) deleted pm_no_index and many "remove hanging casts" --- test/null/test_uops.py | 23 ++++++++++++++++++-- tinygrad/codegen/__init__.py | 9 ++------ tinygrad/uop/ops.py | 41 +++++++++++++++++++++--------------- 3 files changed, 47 insertions(+), 26 deletions(-) diff --git a/test/null/test_uops.py b/test/null/test_uops.py index e730d524b2..f7197fabaf 100644 --- a/test/null/test_uops.py +++ b/test/null/test_uops.py @@ -3,9 +3,9 @@ import unittest import numpy as np from tinygrad.tensor import Tensor from tinygrad.helpers import Timing, Context, cdiv -from tinygrad.dtype import dtypes, ConstFloat, Invalid # noqa: F401 +from tinygrad.dtype import dtypes, AddrSpace, ConstFloat, Invalid # noqa: F401 from tinygrad.device import Device -from tinygrad.uop.ops import Ops, ParamArg, UOp, UPat, dtype_from_uop, exec_alu # noqa: F401 # ParamArg used by eval(str(uop)) roundtrip tests +from tinygrad.uop.ops import Ops, ParamArg, UOp, UPat, dtype_from_uop, exec_alu, graph_rewrite, pm_lower_index_dtype # noqa: F401 # ParamArg used by eval(str(uop)) roundtrip tests from tinygrad.uop.spec import spec_program, spec_shared, type_verify from tinygrad.uop.symbolic import sym from test.helpers import eval_uop, to_uops_list @@ -45,6 +45,25 @@ class TestDTypeFromUOp(unittest.TestCase): with self.assertRaises(RuntimeError): type_verify(UOp.const(weak, value).sink(), spec_program) type_verify(UOp.const(concrete, value).sink(), spec_program) +class TestLowerIndexDtype(unittest.TestCase): + def test_gated_shrink_lowers_to_selected_width(self): + # coalesce builds gated SHRINKs for masked vectorized loads; lowering must resolve them at the + # width the offset bounds select (this one needs long) + buf = UOp.param(0, dtypes.float, (2**31+64,)) + i = UOp.variable("i", 0, 2**28) + shrink = UOp(Ops.SHRINK, src=(buf, (i*24).valid(i < 2**28), UOp.const(dtypes.weakint, 4))) + lowered = graph_rewrite(shrink.sink(), pm_lower_index_dtype) + self.assertTrue(all(u.dtype != dtypes.weakint for u in lowered.backward_slice_with_self), "lowering must resolve all weakint") + sh = next(u for u in lowered.backward_slice_with_self if u.op is Ops.SHRINK) + self.assertEqual(sh.src[1].dtype, dtypes.long) + + def test_reg_buffer_size_lowers(self): + reg = UOp.placeholder((4,), dtypes.float, 0, addrspace=AddrSpace.REG) + self.assertEqual(reg.src[0].dtype, dtypes.weakint) + lowered = graph_rewrite(reg.sink(), pm_lower_index_dtype) + self.assertTrue(all(u.dtype != dtypes.weakint for u in lowered.backward_slice_with_self), "lowering must resolve all weakint") + self.assertEqual(next(u for u in lowered.backward_slice_with_self if u.op is Ops.BUFFER).src[0].dtype, dtypes.int) + class TestSafeCast(unittest.TestCase): def test_cast_folds(self): a = UOp.variable("a", 1, 10, dtype=dtypes.int32) diff --git a/tinygrad/codegen/__init__.py b/tinygrad/codegen/__init__.py index 2b131ecb7d..8da1a2abfa 100644 --- a/tinygrad/codegen/__init__.py +++ b/tinygrad/codegen/__init__.py @@ -38,11 +38,6 @@ pm_number_params = PatternMatcher([ (UPat(Ops.PARAM, name="x"), do_number_param), ]) -pm_no_index = PatternMatcher([ - (UPat(GroupOp.ALU.union({Ops.CONST}), dtype=dtypes.weakint, name="x"), lambda x: x.replace(dtype=dtypes.int)), - (UPat(Ops.CAST, dtype=dtypes.weakint, src=(UPat.var("x"),)), lambda x: x.cast(dtypes.int)), -]) - def build_range_map(sink:UOp) -> dict[int, int]: ctx: dict[int, int] = {} for x in sink.toposort(): @@ -327,7 +322,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: # lower index dtype # NOTE: we need indexing_simplify to remove the cast to long using the Invalid - sink = graph_rewrite(sink, pm_lower_index_dtype+indexing_simplify, name="lower all index dtypes") + sink = graph_rewrite(sink, pm_lower_index_dtype+indexing_simplify, ctx={}, name="lower all index dtypes") # final symbolic before decomp sink = graph_rewrite(sink, symbolic, name="final symbolic") @@ -351,7 +346,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp: # final rules for the renderer (without sym) extra_matcher = ren.extra_matcher if ren.extra_matcher is not None else PatternMatcher([]) - pm_final_rewrite = pm_decomp+extra_matcher+pm_split_ends+pm_no_index + pm_final_rewrite = pm_decomp+extra_matcher+pm_split_ends sink = graph_rewrite(sink, pm_final_rewrite+pm_remove_invalid, ctx=ren, name="final rewrite") # this was the linearizer diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 47de43766b..43569b58c6 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -1699,7 +1699,7 @@ def select_dtype(u:UOp): def lower_alu_dtype(u:UOp, x:UOp, y:UOp, dt:DType) -> UOp: src = u.src[:-2]+(x.cast(dt), y.cast(dt)) return src[0].alu(u.op, *src[1:]).cast(u.dtype) -pm_lower_index_dtype = PatternMatcher([ +pm_lower_weakint = PatternMatcher([ # There are no Unary ops at this point in symbolic, those are introduced later (UPat(Ops.CONST, dtype=dtypes.weakint, name="u"), lambda u: u.replace(dtype=select_dtype(u)).cast(u.dtype) if u.arg!=Invalid else None), # Binary can widen the dtype, WHERE cannot @@ -1707,6 +1707,9 @@ pm_lower_index_dtype = PatternMatcher([ lambda u,x,y: lower_alu_dtype(u, x, y, least_upper_dtype(select_dtype(u), x.dtype, y.dtype))), (UPat(Ops.WHERE, dtypes.weakint, src=(UPat(), UPat.var("x").cast(dtypes.weakint), UPat.var("y").cast(dtypes.weakint)), name="u"), lambda u,x,y: lower_alu_dtype(u, x, y, least_upper_dtype(x.dtype, y.dtype))), + # in a weakint WHERE, an Invalid branch takes the dtype of the other branch + (UPat.var("gate").where(UPat.var("idx", dtypes.ints).cast(dtypes.weakint), UPat(Ops.CONST, arg=Invalid)), + lambda gate,idx: idx.valid(gate).cast(dtypes.weakint)), (UPat(Ops.RANGE, src=(UPat.var("end").cast(dtypes.weakint)), name="r"), lambda r,end: r.replace(dtype=end.dtype, src=(end,)).cast(dtypes.weakint)), (UPat(Ops.STACK, src=UPat().cast(dtypes.weakint), name="v"), lambda v: v.replace(dtype=(dt:=select_dtype(v)), src=tuple(s.src[0].cast(dt) for s in v.src)).cast(dtypes.weakint)), @@ -1717,22 +1720,26 @@ pm_lower_index_dtype = PatternMatcher([ lambda u: u.replace(dtype=None, arg=replace(u.arg, dtype=dtypes.int)).cast(dtypes.weakint) if u.addrspace == AddrSpace.ALU else None), (UPat(Ops.BIND, src=(UPat.var("var").cast(dtypes.weakint), UPat.cvar("val").cast(dtypes.weakint))), lambda var,val: var.bind(val).cast(dtypes.weakint)), - # remove hanging casts - (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx", dtypes.ints).cast()),), lambda buf,idx: buf.index(idx)), - (UPat(Ops.SHRINK, src=(UPat.var("buf"), UPat.var("idx", dtypes.ints).cast(), UPat.var("slen", dtypes.ints).cast(),), name="shrink"), - lambda shrink,buf,idx,slen: shrink.replace(src=(buf,idx,slen))), - (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("gate").where(UPat.var("idx", dtypes.ints).cast(), UPat(Ops.CONST, arg=Invalid)))), - lambda buf,idx,gate: buf.index(idx.valid(gate))), - # remove hanging casts for images - (UPat(Ops.PARAM, src=(UPat.var("shape").cast(),), name="p"), lambda p,shape: p.replace(src=(shape,))), - (UPat(Ops.INDEX, src=(UPat.var("buf"), UPat.var("idx_y", dtypes.ints).cast(), UPat.var("idx_x", dtypes.ints).cast()),), - lambda buf,idx_x,idx_y: buf.index(idx_y, idx_x, dtype=dtypes.float)), - (UPat(Ops.INDEX, src=(UPat.var("buf"), - UPat.var("gate").where(UPat.var("idx_y", dtypes.ints).cast(), UPat(Ops.CONST, arg=Invalid)), - UPat.var("gate").where(UPat.var("idx_x", dtypes.ints).cast(), UPat(Ops.CONST, arg=Invalid)))), - lambda buf,idx_x,idx_y,gate: buf.index(idx_y.valid(gate), idx_x.valid(gate), dtype=dtypes.float)), - (UPat((Ops.SINK, Ops.NOOP, Ops.END), name="n"), - lambda n: n.replace(src=tuple(s.src[0] if s.op is Ops.CAST and s.dtype == dtypes.weakint else s for s in n.src))), +]) +def lower_weak_srcs(ctx:dict[UOp, UOp]|None, u:UOp) -> UOp|None: + if ctx is None: ctx = {} + def lower(s:UOp) -> UOp: + if (r:=ctx.get(s)) is None: + r = graph_rewrite(s, pm_lower_weakint) + # the consumer absorbs the cast on its own edge + ctx[s] = r = r.src[0] if r.op is Ops.CAST and r.dtype == dtypes.weakint else r + return r + # a comparison demands a common operand width: lower it whole so the Binary rule unifies its operands + ret = lower(u) if u.op in GroupOp.Comparison else u.replace(src=tuple(lower(s) if s.dtype == dtypes.weakint else s for s in u.src)) + return None if ret is u else ret +pm_lower_index_dtype = PatternMatcher([ + (UPat(GroupOp.All, name="u"), + lambda ctx,u: lower_weak_srcs(ctx, u) if u.dtype != dtypes.weakint and any(s.dtype == dtypes.weakint for s in u.src) else None), + # a valid index into an n-element buffer lives in [0,n): a gated long index narrows when n-1 fits int32 (out-of-gate wraps, discarded) + # TODO: more generic + (UPat((Ops.INDEX, Ops.SHRINK), src=(UPat.var("buf"), UPat.var("gate").where(UPat.var("idx", dtypes.long), UPat(Ops.CONST, arg=Invalid))), + allow_any_len=True, name="u"), + lambda u,buf,gate,idx: u.replace(src=(buf, idx.cast(dtypes.int).valid(gate))+u.src[2:]) if buf.max_numel()-1 <= dtypes.int32.max else None), ]) def _index_to_concrete_int(u:UOp) -> UOp: return graph_rewrite(u.sink(), pm_lower_index_dtype).src[0]