mirror of https://github.com/commaai/tinygrad.git
136 lines
6.7 KiB
Python
136 lines
6.7 KiB
Python
from __future__ import annotations
|
|
from typing import Dict, Union, Tuple, Any, List, cast
|
|
import functools, hashlib
|
|
from enum import Enum, auto
|
|
from dataclasses import dataclass
|
|
from tinygrad.helpers import dedup, pretty_print, prod
|
|
from tinygrad.ops import ReduceOps, UnaryOps, BinaryOps, TernaryOps, UOp, UOps
|
|
from tinygrad.dtype import ImageDType, PtrDType, dtypes, DType, ConstType
|
|
from tinygrad.shape.symbolic import Variable, sint
|
|
from tinygrad.shape.shapetracker import ShapeTracker
|
|
|
|
# these ops are deleted after AST is UOp
|
|
class BufferOps(Enum): LOAD = auto(); CONST = auto(); STORE = auto() # noqa: E702
|
|
class MetaOps(Enum): KERNEL = auto();
|
|
Op = Union[UnaryOps, BinaryOps, ReduceOps, MetaOps, TernaryOps, BufferOps]
|
|
|
|
@dataclass(frozen=True)
|
|
class MemBuffer:
|
|
idx: int
|
|
dtype: DType
|
|
st: ShapeTracker
|
|
|
|
@dataclass(frozen=True)
|
|
class ConstBuffer:
|
|
val: ConstType | Variable
|
|
dtype: DType
|
|
st: ShapeTracker
|
|
|
|
@dataclass(frozen=True, eq=False)
|
|
class LazyOp:
|
|
op: Op
|
|
src: Tuple[LazyOp, ...] = ()
|
|
arg: Any = None
|
|
def cached_compare(self, x, context):
|
|
if id(self) == id(x): return True
|
|
if self.op != x.op or self.arg != x.arg or len(self.src) != len(x.src): return False
|
|
if (key := (id(self), id(x))) in context: return context[key]
|
|
ret = context[key] = all(a.cached_compare(b, context) for a,b in zip(self.src, x.src))
|
|
return ret
|
|
def __eq__(self, x): return self.cached_compare(x, context={})
|
|
def __repr__(self:LazyOp): return pretty_print(self, lambda x: f'LazyOp({x.op}, arg={x.arg}, src=(%s))')
|
|
@functools.cached_property
|
|
def dtype(self) -> DType:
|
|
if self.op in BufferOps: return self.arg.dtype
|
|
if self.op in [UnaryOps.CAST, UnaryOps.BITCAST]: return self.arg
|
|
return dtypes.bool if self.op in {BinaryOps.CMPLT, BinaryOps.CMPNE} else self.src[-1].dtype
|
|
@functools.cached_property
|
|
def full_shape(self) -> Tuple[sint, ...]:
|
|
if len(self.src) == 0 and self.op in BufferOps: return self.arg.st.shape
|
|
return tuple(max(x) for x in zip(*[x.full_shape for x in self.src]))
|
|
@functools.cached_property
|
|
def key(self) -> bytes:
|
|
return hashlib.sha256(functools.reduce(lambda x,y: x+y, [s.key for s in self.src], str((self.op, self.arg)).encode())).digest()
|
|
@functools.cached_property
|
|
def hash(self): return hash((self.op, self.src, self.arg))
|
|
def __hash__(self): return self.hash
|
|
@functools.cached_property
|
|
def lazyops(self) -> List[LazyOp]: return dedup([self] + [item for x in self.src for item in x.lazyops])
|
|
def vars(self) -> List[Variable]:
|
|
extract_vars = [x.arg.st.vars() for x in self.lazyops if x.op in BufferOps]
|
|
const_vars = [x.arg.val for x in self.lazyops if x.op is BufferOps.CONST and isinstance(x.arg.val, Variable)]
|
|
return sorted(set.union(*extract_vars, set(const_vars)), key=lambda v: v.expr)
|
|
def __add__(self, x:LazyOp): return LazyOp(BinaryOps.ADD, (self, x))
|
|
def __sub__(self, x:LazyOp): return LazyOp(BinaryOps.ADD, (self, -x))
|
|
def __mul__(self, x:LazyOp): return LazyOp(BinaryOps.MUL, (self, x))
|
|
def ne(self, x:LazyOp): return LazyOp(BinaryOps.CMPNE, (self, x))
|
|
def eq(self, x:LazyOp): return -self.ne(x)
|
|
def __neg__(self): return LazyOp(UnaryOps.NEG, (self,))
|
|
@staticmethod
|
|
def const(val, dtype:DType, shape:Tuple[sint, ...]):
|
|
return LazyOp(BufferOps.CONST, (), ConstBuffer(val, dtype, ShapeTracker.from_shape(()).reshape((1,)*len(shape)).expand(shape)))
|
|
|
|
# the living definition of LazyOps
|
|
def verify_lazyop(ast:LazyOp) -> Dict[LazyOp, ShapeTracker]:
|
|
assert ast.op is MetaOps.KERNEL, "must be SINK"
|
|
sts: Dict[LazyOp, ShapeTracker] = {}
|
|
def assert_valid(op:LazyOp, st:ShapeTracker):
|
|
if op in sts: return
|
|
# restore globals from the two stage reduce
|
|
if op.op is BufferOps.LOAD and op.arg.idx < 0:
|
|
assert_valid(local_reduce:=op.src[0].src[0], op.arg.st)
|
|
return sts.setdefault(op, sts[local_reduce])
|
|
for x in op.src: assert_valid(x, st)
|
|
# only reduceop is allowed to change shape, limited to turning n to 1
|
|
if op.op in ReduceOps:
|
|
axis = op.arg
|
|
assert isinstance(axis, tuple) and all(isinstance(i, int) for i in axis), f"reduceop must have axis {op.arg}"
|
|
st = ShapeTracker.from_shape(sts[op.src[0]].reduce(axis))
|
|
else:
|
|
# movementops are pushed to the edges with LOAD
|
|
# elementwise inherits shape
|
|
st = op.arg.st if op.op in BufferOps else sts[op.src[0]]
|
|
for x in op.src:
|
|
if sts[x].shape != st.shape:
|
|
if prod(sts[x].shape) == prod(st.shape): raise AssertionError(f"found implicit reshape {x.op} {op.op} {sts[x].shape} != {st.shape}")
|
|
raise AssertionError(f"found implicit expand {x.op} {sts[x].shape} != {op.op} {st.shape} {prod(sts[x].shape)} != {prod(st.shape)}")
|
|
sts[op] = st
|
|
for i, out in enumerate(ast.src):
|
|
assert out.arg.idx == i, f"unexpected output buffer idx {out.arg.idx} != {i}"
|
|
assert out.op is BufferOps.STORE, f"kernels must have stores as the output, got {out.op}"
|
|
assert out.arg.st.size == ast.src[-1].arg.st.size, f"outputs must have the same size, got {out.arg.st.size}"
|
|
assert_valid(out, out.arg.st)
|
|
shape_dims = [sorted(dedup(dims)) for dims in zip(*[x.shape for x in sts.values()])]
|
|
assert all(len(x) == 1 or (len(x) == 2 and x[0] == 1) for x in shape_dims), f"shapes must have either 1 or n in each dimension, {shape_dims}"
|
|
return sts
|
|
|
|
def to_uop(*a) -> UOp:
|
|
assert isinstance(a[0], LazyOp), f"{a} must be a LazyOp ast"
|
|
if a[0].op is BufferOps.STORE: ast = LazyOp(MetaOps.KERNEL, a)
|
|
else:
|
|
assert a[0].op is MetaOps.KERNEL
|
|
ast = a[0]
|
|
verify_lazyop(ast)
|
|
@functools.lru_cache(None)
|
|
def create_uop(lop:LazyOp) -> UOp:
|
|
if lop.op in BufferOps:
|
|
st_uop = lop.arg.st.to_uop()
|
|
membuf_dtype: DType = lop.arg.dtype
|
|
dtype = membuf_dtype.base if isinstance(membuf_dtype, ImageDType) else membuf_dtype
|
|
if lop.op is BufferOps.CONST:
|
|
return UOp(UOps.CONST, dtype, (st_uop,), lop.arg.val)
|
|
buf = UOp(UOps.DEFINE_GLOBAL, membuf_dtype if isinstance(membuf_dtype, ImageDType) else PtrDType(membuf_dtype), (), lop.arg.idx)
|
|
if lop.op is BufferOps.LOAD: return UOp(UOps.LOAD, dtype, (buf, st_uop))
|
|
return UOp(UOps.STORE, dtypes.void, (buf, st_uop, create_uop(lop.src[0])))
|
|
src = tuple(create_uop(x) for x in lop.src)
|
|
if lop.op is MetaOps.KERNEL: return UOp(UOps.SINK, dtypes.void, src)
|
|
if lop.op in ReduceOps:
|
|
alu_op = {ReduceOps.SUM:BinaryOps.ADD, ReduceOps.PROD:BinaryOps.MUL, ReduceOps.MAX:BinaryOps.MAX}[cast(ReduceOps, lop.op)]
|
|
return UOp(UOps.REDUCE_AXIS, src[0].dtype, src, (alu_op, lop.arg))
|
|
if lop.op is UnaryOps.CAST: return UOp(UOps.CAST, lop.arg.scalar(), src)
|
|
if lop.op is UnaryOps.BITCAST: return UOp(UOps.BITCAST, lop.arg.scalar(), src)
|
|
return src[0].alu(lop.op, *src[1:])
|
|
ret = create_uop(ast)
|
|
#with open("/tmp/ast", "w") as f: f.write(str(ret))
|
|
return ret
|