From b252f890da4246f2cce6eceeec7badb4361af30c Mon Sep 17 00:00:00 2001 From: George Hotz <72895+geohot@users.noreply.github.com> Date: Mon, 29 Sep 2025 14:55:01 +1000 Subject: [PATCH] add support for SPEC=1 (#12322) * add support for SPEC=1 * cleaner place for it * non rangeify spec * split non rangeify --- tinygrad/helpers.py | 1 + tinygrad/uop/ops.py | 11 +++++--- tinygrad/uop/spec.py | 60 ++++++++++++++++++++++++++++++++++++++++---- tinygrad/uop/upat.py | 3 ++- 4 files changed, 66 insertions(+), 9 deletions(-) diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index 95628c300a..a0fa11584a 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -147,6 +147,7 @@ EMULATE = ContextVar("EMULATE", "") CPU_COUNT = ContextVar("CPU_COUNT", max(1, (os.cpu_count() or 1) // (4 if ARCH_X86 else 2))) # take 1/2 of the cores, accounting HT CPU_LLVM, AMD_LLVM = ContextVar("CPU_LLVM", 0), ContextVar("AMD_LLVM", 1) VIZ = PROFILE = ContextVar("VIZ", 0) +SPEC = ContextVar("SPEC", 0) @dataclass(frozen=True) class Metadata: diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py index 8bd4f80d48..206d0bdf41 100644 --- a/tinygrad/uop/ops.py +++ b/tinygrad/uop/ops.py @@ -7,7 +7,7 @@ from tinygrad.uop import Ops, GroupOp from tinygrad.uop.mathtraits import MathTrait from tinygrad.dtype import ConstType, ImageDType, dtypes, DType, truncate, PtrDType, least_upper_dtype, Invalid, InvalidType from tinygrad.helpers import ContextVar, all_int, prod, getenv, all_same, Context, partition, temp, unwrap, T, argfix, Metadata, flatten, TRACEMETA -from tinygrad.helpers import PICKLE_BUFFERS, PROFILE, dedup, cdiv, cmod, diskcache_put, to_function_name, cpu_profile, TracingKey, RANGEIFY, VIZ +from tinygrad.helpers import PICKLE_BUFFERS, PROFILE, dedup, cdiv, cmod, diskcache_put, to_function_name, cpu_profile, TracingKey, RANGEIFY, VIZ, SPEC if TYPE_CHECKING: from tinygrad.shape.shapetracker import ShapeTracker from tinygrad.device import Buffer, MultiBuffer @@ -66,6 +66,10 @@ class UOpMetaClass(type): if _buffer is not None: assert op is Ops.BUFFER, f"trying to set Buffer {_buffer} for {op}" buffers[created] = _buffer + if SPEC: + from tinygrad.uop.spec import full_spec + with Context(IGNORE_OOB=1): ret = full_spec.rewrite(created) + if cast(bool|None, ret) is not True: raise RuntimeError(f"SPEC ISSUE {ret}: {created}") return created # some uops map to other stuff @@ -624,7 +628,7 @@ class UOp(MathTrait, metaclass=UOpMetaClass): return fxn(**{k:v for k,v in var_vals.items() if k in varnames}) def render(self, simplify=True, pm:PatternMatcher|None=None) -> str: - with Context(TRACK_MATCH_STATS=0): + with Context(TRACK_MATCH_STATS=0, SPEC=0): ret = graph_rewrite(self.simplify() if simplify else self, renderer if pm is None else pm) return ret.arg if ret.op is Ops.NOOP else str(ret) @@ -1072,7 +1076,7 @@ pm_lower_index_dtype = PatternMatcher([ (UPat(Ops.WHERE, dtypes.index, src=(UPat(), UPat.var("x"), UPat(Ops.CONST, arg=Invalid)), name="u"), lambda u,x: u.replace(dtype=x.dtype)), (UPat(Ops.WHERE, dtypes.index, src=(UPat.var("cond"), UPat.var("x"), UPat.var("y"))), lambda cond,x,y: cond.where(x.cast(dt:=least_upper_dtype(x.dtype, y.dtype)), y.cast(dt))), - (UPat((Ops.CONST, Ops.VCONST), dtype=dtypes.index, name="u"), lambda u: u.replace(dtype=select_dtype(u))), + (UPat((Ops.CONST, Ops.VCONST), dtype=dtypes.index, name="u"), lambda u: u.replace(dtype=select_dtype(u)) if u.arg != Invalid else None), (UPat((Ops.RANGE,), dtype=dtypes.index, src=(UPat.var("end")), name="r"), lambda ctx,r,end: r.replace(dtype=(dt:=select_dtype(r)), src=(end.cast(dt),))), (UPat(Ops.CAST, dtype=dtypes.index, src=(UPat.var("x", dtypes.ints),), name="u"), lambda u,x: x), @@ -1137,6 +1141,7 @@ pm_pyrender = PatternMatcher([ lambda x: UOp(Ops.NOOP, arg=f"{x.src[0].arg}.f({x.op}, dtype=dtypes.bool)")), ]) +@Context(SPEC=0) def pyrender(ast:UOp) -> list[str]: cmap = ast.get_children_map() to_render = set() diff --git a/tinygrad/uop/spec.py b/tinygrad/uop/spec.py index 48737fc9e2..04457b6b3b 100644 --- a/tinygrad/uop/spec.py +++ b/tinygrad/uop/spec.py @@ -1,7 +1,7 @@ from typing import cast, Callable from tinygrad.uop.ops import PatternMatcher, UPat, GroupOp, Ops, UOp, print_uops, python_alu, graph_rewrite, AxisType from tinygrad.dtype import DType, ImageDType, dtypes, PtrDType, AddrSpace, Invalid -from tinygrad.helpers import all_same, prod, DEBUG, ContextVar, Context, cpu_profile +from tinygrad.helpers import all_same, prod, DEBUG, ContextVar, Context, cpu_profile, RANGEIFY from tinygrad.shape.shapetracker import ShapeTracker try: import z3 @@ -49,7 +49,7 @@ try: ]) def uops_to_z3(solver, *uops: UOp) -> 'list[z3.ExprRef]': - with Context(TRACK_MATCH_STATS=0): # cant pickle z3 objects + with Context(TRACK_MATCH_STATS=0, SPEC=0): # cant pickle z3 objects, and these UOps don't follow spec return [s.arg[1] for s in graph_rewrite(uops[0].sink(*uops[1:]), z3_renderer, ctx=(solver, {})).src] z3_imported = True @@ -124,7 +124,8 @@ tensor_uop_spec = buffer_spec+assign_spec+PatternMatcher([ # ***** uop type spec ***** -def validate_index(idx:UOp, gate:UOp=UOp.const(dtypes.bool, True)): +def validate_index(idx:UOp, gate:UOp|None=None): + if gate is None: gate = UOp.const(dtypes.bool, True) # TODO: check for overflow if IGNORE_OOB or isinstance(idx.dtype, ImageDType) or (sz := idx.src[0].ptrdtype.size) == -1: return True # We can use UOp min/max to do a faster check, but it can give false positive since its not an exact bound and doesn't consider the mask @@ -146,7 +147,8 @@ def validate_index(idx:UOp, gate:UOp=UOp.const(dtypes.bool, True)): return False return True -def validate_store(idx:UOp, val:UOp, gate:UOp=UOp.const(dtypes.bool, True)): +def validate_store(idx:UOp, val:UOp, gate:UOp|None=None): + if gate is None: gate = UOp.const(dtypes.bool, True) if gate.op is Ops.IF: gate = gate.src[0] # we need to find the implicit gates, inverse of delete_redundant_gates for u in val.toposort(): @@ -226,7 +228,7 @@ spec = PatternMatcher([ (UPat(Ops.REDUCE_AXIS, name="x"), lambda x: isinstance(x.arg, tuple) and len(x.arg) >= 2 and x.arg[0] in {Ops.ADD, Ops.MUL, Ops.MAX}), (UPat(Ops.GEP, src=(UPat.var("src"),), name="gep"), lambda gep,src: gep.dtype == src.dtype.scalar()), - (UPat(Ops.VECTORIZE, name="x"), lambda x: len(x.src)>1 and len(x.src) == x.dtype.count and all(x.dtype == y.dtype.vec(len(x.src)) for y in x.src)), + (UPat(Ops.VECTORIZE, name="x"), lambda x: len(x.src)>1 and len(x.src) == x.dtype.vcount and all(x.dtype == y.dtype.vec(len(x.src)) for y in x.src)), (UPat((Ops.BITCAST, Ops.CAST), src=(UPat(),), name="x"), lambda x: x.arg is None), (UPat(Ops.BARRIER, dtypes.void, src=UPat(Ops.STORE, allow_any_len=True)), lambda: True), # NOTE: all pointers must be local (UPat(Ops.BARRIER, dtypes.void), lambda: True), # BARRIERs can also happen at the end of loops @@ -250,6 +252,54 @@ ast_spec = PatternMatcher([ (UPat(GroupOp.All-{Ops.SINK}, name="root"), lambda root: all_same([x.shape for x in root.src if x.st is not None])), ]) +# *** this spec should match all UOps ever created *** + +full_non_rangeify_spec = PatternMatcher([]) if RANGEIFY else PatternMatcher([ + # in non rangeify const can still have a View, and sometimes a FUSE while propagating + (UPat((Ops.VIEW, Ops.FUSE)).f(Ops.CONST), lambda: True), +]) + +full_spec = PatternMatcher([ + # Invalid must have type Index + (UPat(Ops.CONST, arg=Invalid, name="x"), lambda x: x.dtype.scalar() == dtypes.index), + # where on index in rhs position is fine + (UPat(Ops.WHERE, src=(UPat(dtype=dtypes.bool), UPat(), UPat(dtype=dtypes.index))), lambda: True), + + # all children is fine + (UPat(Ops.CHILDREN), lambda: True), + # child must have CHILDREN parent + (UPat(Ops.CHILD, src=(UPat(Ops.CHILDREN),)), lambda: True), + + # all rewrite error are okay + (UPat(Ops.REWRITE_ERROR), lambda: True), + + # buffer view with index or load is okay + (UPat(Ops.BUFFER_VIEW, src=(UPat((Ops.INDEX, Ops.LOAD)),)), lambda: True), + + # linearizer outputs + intermediate KERNELs + (UPat((Ops.BLOCKSTART, Ops.BLOCK, Ops.BLOCKFINAL, Ops.BLOCKEND, Ops.KERNEL), dtype=dtypes.void), lambda: True), + # realize is fine + (UPat((Ops.REALIZE, Ops.BUFFERIZE)), lambda: True), + # expander: unroll/contract/gep + (UPat((Ops.UNROLL, Ops.CONTRACT, Ops.GEP, Ops.CAT, Ops.PTRCAT)), lambda: True), + # any vectorize is okay? + (UPat(Ops.VECTORIZE), lambda: True), + # index/copy op during RANGEIFY + (UPat((Ops.INDEX, Ops.COPY, Ops.REDUCE)), lambda: True), + # all loads/stores + (UPat((Ops.LOAD, Ops.STORE)), lambda: True), + # all ifs + (UPat(Ops.IF), lambda: True), + # all assign + (UPat(Ops.ASSIGN), lambda: True), + # all DEFINE_VAR to deal with the floats used in reduce collapse + (UPat(Ops.DEFINE_VAR), lambda: True), + # allow index type + (UPat(GroupOp.All, dtype=dtypes.index), lambda: True), + # reshape on STORE + (UPat(Ops.RESHAPE, src=(UPat(Ops.STORE),)), lambda: True), +])+full_non_rangeify_spec+tensor_uop_spec+spec + # ***** uop helpers ***** def type_verify(uops:list[UOp], extra_spec:PatternMatcher|None=None): diff --git a/tinygrad/uop/upat.py b/tinygrad/uop/upat.py index 5c5f221d31..32255ec1fa 100644 --- a/tinygrad/uop/upat.py +++ b/tinygrad/uop/upat.py @@ -153,7 +153,8 @@ def _get_code(self:UPat, has_ctx:bool): @functools.cache def upat_compile(self:UPat, fxn) -> Callable|None: real_fxn = types.FunctionType(*deconstruct_function(fxn)) - code = _get_code(self, 'ctx' in inspect.signature(real_fxn).parameters) + # UOps used here don't follow the spec + with Context(SPEC=0): code = _get_code(self, 'ctx' in inspect.signature(real_fxn).parameters) if code is None: return None code_str, dyn_lookup = code globs = dyn_lookup.copy()