forked from tinygrad/tinygrad
rework pm_lower_weakint [pr] (#17136)
deleted pm_no_index and many "remove hanging casts"
This commit is contained in:
chenyu
+21
-2
@@ -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)
|
||||
|
||||
@@ -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
|
||||
|
||||
+24
-17
@@ -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]
|
||||
|
||||
|
||||
Reference in New Issue
Block a user