forked from tinygrad/tinygrad
add support for SPEC=1 (#12322)
* add support for SPEC=1 * cleaner place for it * non rangeify spec * split non rangeify
This commit is contained in:
George Hotz
@@ -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:
|
||||
|
||||
+8
-3
@@ -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()
|
||||
|
||||
+55
-5
@@ -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):
|
||||
|
||||
@@ -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()
|
||||
|
||||
Reference in New Issue
Block a user