diff --git a/.gitignore b/.gitignore index 4911dd7c18..25faa75a96 100644 --- a/.gitignore +++ b/.gitignore @@ -3,6 +3,7 @@ notebooks .*.swp .*.swo *.pyc +*.so build dist *.egg-info diff --git a/compile.sh b/compile.sh new file mode 100755 index 0000000000..93490b1505 --- /dev/null +++ b/compile.sh @@ -0,0 +1,5 @@ +#!/bin/bash +mypyc tinygrad/llops/ops_gpu.py tinygrad/shape/__init__.py tinygrad/ops.py tinygrad/ast.py \ + tinygrad/helpers.py tinygrad/mlops.py tinygrad/nn/__init__.py tinygrad/graph.py tinygrad/lazy.py \ + tinygrad/tensor.py + diff --git a/examples/benchmark_train_efficientnet.py b/examples/benchmark_train_efficientnet.py index 1d8208e749..9255679903 100644 --- a/examples/benchmark_train_efficientnet.py +++ b/examples/benchmark_train_efficientnet.py @@ -66,7 +66,7 @@ if __name__ == "__main__": CL.CACHE = None mem_used = CL.mem_used - loss_cpu = loss.detach().cpu().data[0] + loss_cpu = loss.detach().numpy()[0] cl = time.monotonic() print(f"{(st-cpy)*1000.0:7.2f} ms cpy, {(cl-st)*1000.0:7.2f} ms run, {(mt-st)*1000.0:7.2f} ms build, {(et-mt)*1000.0:7.2f} ms realize, {(cl-et)*1000.0:7.2f} ms CL, {loss_cpu:7.2f} loss, {tensors_allocated():4d} tensors, {mem_used/1e9:.2f} GB used, {ops*1e-9/(cl-st):9.2f} GFLOPS") diff --git a/extra/lib_test_ast.py b/extra/lib_test_ast.py index 549a693590..e41244d773 100644 --- a/extra/lib_test_ast.py +++ b/extra/lib_test_ast.py @@ -3,8 +3,7 @@ import numpy as np from typing import Dict, Type from tinygrad.ast import ASTKernel from tinygrad.llops.ops_cpu import CPUBuffer -from tinygrad.ops import DeviceBuffer -from tinygrad.lazy import map_buffers +from tinygrad.ops import DeviceBuffer, map_buffers in_test = False test_cnt = 0 diff --git a/rmso.sh b/rmso.sh new file mode 100755 index 0000000000..46817947ff --- /dev/null +++ b/rmso.sh @@ -0,0 +1,3 @@ +#!/bin/bash +rm tinygrad/*.so tinygrad/shape/*.so tinygrad/llops/*.so tinygrad/nn/*.so tinygrad/runtime/*.so *.so + diff --git a/tinygrad/ast.py b/tinygrad/ast.py index 4a880e41a7..15b839336c 100644 --- a/tinygrad/ast.py +++ b/tinygrad/ast.py @@ -1,6 +1,6 @@ -from enum import Enum +from enum import Enum, auto import itertools -from typing import List, Tuple +from typing import List, Tuple, Optional from tinygrad.helpers import prod, dedup, all_same from tinygrad.ops import LazyOp, MovementOps, get_lazyop_info, get_buffers, ReduceOps, get_lazyops from tinygrad.shape import ShapeTracker, View, strides_for_shape @@ -11,14 +11,15 @@ def get_first_reduce(shapes): return i return len(shapes[0]) # off the end -Types = Enum("Types", ["FLOAT", "FLOAT4"]) +# this will be removed soon anyway +class Types(Enum): FLOAT = auto(); FLOAT4 = auto() # noqa: E702 class Token: def __init__(self, tok:str, typ:Types, ptr:bool=False): assert isinstance(tok, str) self.tok, self.typ, self.ptr = tok, typ, ptr self.axis : List[Tuple[int, int, bool]] = [] def array(self, length, stride, reduce): self.axis.append((length, stride, reduce)) - def size(self): return prod(x[0] for x in self.axis) + def size(self): return prod([x[0] for x in self.axis]) def offsets(self): return [sum(t) for t in itertools.product(*[[y*x[1] for y in range(x[0])] for x in self.axis[::-1]])] if len(self.axis) else [0] # TODO: this is sort of a hack, it gets the accumulator indices def acc_offsets(self): @@ -53,12 +54,13 @@ class ASTKernel: # TODO: should be optional if it's hitting a function cache self.processed = False - def process(self): + def process(self) -> None: if self.processed: return self.processed = True reduceops = [x for x in get_lazyops(self.ast) if x.op in ReduceOps] assert len(dedup(reduceops)) <= 1, "max one reduce op in an ast" self.reduceop = reduceops[0] if reduceops else None + self.reduceopop : Optional[ReduceOps] = self.reduceop.op if self.reduceop is not None and isinstance(self.reduceop.op, ReduceOps) else None self.earlybufs = dedup(get_buffers(self.reduceop)) if self.reduceop else [] self.buftokens = [Token(f"data{i}", Types.FLOAT, ptr=True) for i in range(len(self.bufs))] @@ -98,7 +100,7 @@ class ASTKernel: print_ast(self.input_ast, "ast") @property - def shape_len(self): return len(self.sts[0].shape) + def shape_len(self) -> int: return len(self.sts[0].shape) def simplify_ones(self): # remove places where the shape is all ones diff --git a/tinygrad/lazy.py b/tinygrad/lazy.py index 56f0f2d453..5586fcf9fa 100644 --- a/tinygrad/lazy.py +++ b/tinygrad/lazy.py @@ -1,9 +1,10 @@ from __future__ import annotations -from typing import Optional, Tuple, Union, List, Dict, Any, Final +from typing import Optional, Tuple, Union, List, Dict, Any, ClassVar import sys, weakref +from weakref import WeakValueDictionary from tinygrad.helpers import ConvArgs, prod from tinygrad.shape import ShapeTracker -from tinygrad.ops import DeviceBuffer, UnaryOps, BinaryOps, ReduceOps, MovementOps, ProcessingOps, LoadOps, OpType, LazyOp, get_buffers, DEBUG +from tinygrad.ops import DeviceBuffer, UnaryOps, BinaryOps, ReduceOps, MovementOps, ProcessingOps, LoadOps, OpType, LazyOp, get_buffers, map_buffers, DEBUG from tinygrad.graph import log_op from tinygrad.helpers import getenv @@ -22,12 +23,6 @@ REMOVE_MOVEMENT_NOPS, MERGE_UNARY_OPS, MERGE_ELEMENTWISE_INTO_REDUCE, SHUFFLE_MO MERGE_ELEMENTWISE_OPS, MERGE_ONE_REDUCE_INTO_ELEMENTWISE = OPT>=2, OPT>=2 SHUFFLE_PAD_OPS = OPT>=3 # NOTE: 0/0 is NaN if you pad, so this can change the output -# **** realize helpers **** -def map_buffers(real_srcs, x:LazyOp) -> LazyOp: - if x in real_srcs: - return map_buffers(real_srcs, real_srcs[x]) if isinstance(real_srcs[x], LazyOp) else real_srcs[x] - return LazyOp(x.op, tuple(map_buffers(real_srcs, y) for y in x.src), x.arg) - # **** realize functions **** def _ast_reduceops(self:LazyBuffer) -> LazyOp: # TODO: this can also corealize a binary op after the reduce, not just before @@ -74,18 +69,25 @@ def get_weakop(op:LazyOp) -> LazyOp: return LazyOp(op.op, tuple(get_weakop(x) if def get_movementroot(root:LazyBuffer) -> LazyBuffer: return get_movementroot(root.op.src[0]) if root.realized is None and (root.optype == MovementOps or (root.op.op == LoadOps.CONTIGUOUS and root.op.src[0].st.contiguous)) else root def get_movementroot_contiguous(x:LazyBuffer) -> LazyBuffer: return get_movementroot(x) if x.optype == MovementOps and x.st.contiguous else x +def replace_with_movement_op(y:Union[LazyOp, LazyBuffer], op:MovementOps, arg:Tuple[Any, ...]) -> LazyBuffer: + if isinstance(y, LazyBuffer): return y.movement_op(op, arg) + assert y.op in BinaryOps or y.op in UnaryOps + return elementwise_op(y.op, *[replace_with_movement_op(z, op, arg) for z in y.src]) # type: ignore + +def support_weakref(x): return x +@support_weakref # needed for mypyc, this prevents LazyBuffer from becoming a native class class LazyBuffer: __deletable__ = ('op',) - lazycache : Final[weakref.WeakValueDictionary[Tuple[str, OpType, LazyOp], LazyBuffer]] = weakref.WeakValueDictionary() + lazycache : ClassVar[WeakValueDictionary[Tuple[str, OpType, LazyOp], LazyBuffer]] = WeakValueDictionary() def __new__(cls, device:str, shape:Union[ShapeTracker, Tuple[int, ...]], optype:OpType, op:LazyOp): # fromcpu aren't cached if optype == LoadOps and op.op == LoadOps.FROMCPU: return super().__new__(cls) wop = (device, optype, get_weakop(op)) # NOTE: shape should be deterministic. annoying to cache with the ShapeTracker # NOTE: we need "ret" to prevent the new buffer from being immediately deleted - if wop not in LazyBuffer.lazycache: - LazyBuffer.lazycache[wop] = ret = super().__new__(cls) # noqa: F841, pylint: disable=W0612 - return LazyBuffer.lazycache[wop] + if wop not in LazyBuffer.lazycache: LazyBuffer.lazycache[wop] = ret = super().__new__(cls) + else: ret = LazyBuffer.lazycache[wop] + return ret def __init__(self, device:str, shape:Union[ShapeTracker, Tuple[int, ...]], optype:OpType, op:LazyOp): if hasattr(self, 'device'): @@ -140,7 +142,7 @@ class LazyBuffer: self.realized = self.dbuffer.exec_ast(map_buffers({x:x.realize(self.device) for x in get_buffers(ast)}, ast)) log_op(self.realized, ast) - assert self.realized.shape == self.shape + assert self.realized.shape == self.shape, f"shape mismatch on realize {self.realized.shape} vs {self.shape}" assert isinstance(self.realized, Device._buffers[self.device]) return self.realized @@ -168,7 +170,7 @@ class LazyBuffer: padding = tuple((max(0, -p[0]), max(0, p[1]-self.shape[i])) for i,p in enumerate(arg)) return self.movement_op(MovementOps.PAD, padding).movement_op(MovementOps.SHRINK, tuple((p[0] + padding[i][0], p[1] + padding[i][0]) for i,p in enumerate(arg))) - def movement_op(self:LazyBuffer, op:MovementOps, arg : Tuple[Any, ...]) -> LazyBuffer: + def movement_op(self:LazyBuffer, op:MovementOps, arg:Tuple[Any, ...]) -> LazyBuffer: # very instant nop if op == MovementOps.RESHAPE and self.shape == arg: return self @@ -200,12 +202,7 @@ class LazyBuffer: # if this MovementOp is being applied to a BinaryOp, apply the MovementOp to all the BinaryOp inputs instead if SHUFFLE_MOVEMENT_OPS and self.optype == BinaryOps and self.realized is None and len(self.children) == 0 and (SHUFFLE_PAD_OPS or op != MovementOps.PAD) and op not in [MovementOps.EXPAND, MovementOps.STRIDED]: - def replace_with_movement_op(y:Union[LazyOp, LazyBuffer]) -> LazyBuffer: - if isinstance(y, LazyBuffer): - return y.movement_op(op, arg) - assert y.op in BinaryOps or y.op in UnaryOps - return elementwise_op(y.op, *[replace_with_movement_op(z) for z in y.src]) # type: ignore - return replace_with_movement_op(self.op) + return replace_with_movement_op(self.op, op, arg) # create the buffer ret = LazyBuffer(self.device, ShapeTracker(self.st).movement_op(op, arg), MovementOps, LazyOp(op, (self,), arg)) diff --git a/tinygrad/llops/ops_gpu.py b/tinygrad/llops/ops_gpu.py index 692d4fb74e..39461433dd 100644 --- a/tinygrad/llops/ops_gpu.py +++ b/tinygrad/llops/ops_gpu.py @@ -1,6 +1,6 @@ from __future__ import annotations import numpy as np -from typing import List, Tuple, Optional, Dict, Union, Set +from typing import List, Tuple, Optional, Dict, Union, Set, Final, Callable from tinygrad.helpers import prod from tinygrad.ops import DEBUG, UnaryOps, BinaryOps, ReduceOps, MovementOps, LazyOp, Op, ExplicitExecAST, GlobalCounters from tinygrad.ast import ASTKernel, Token, Types @@ -10,7 +10,7 @@ from tinygrad.shape.symbolic import ModNode # this will go away when VALIDHACK from tinygrad.helpers import getenv CUDA = getenv("CUDA", 0) -if not CUDA: from tinygrad.runtime.opencl import CLBuffer, CLImage, CLProgram, CL # NOTE: using CL will not work for the CUDA runtime # noqa: F401 +if not CUDA: from tinygrad.runtime.opencl import CLBuffer, CLImage, CLProgram, CL, cl as pyopencl # NOTE: using CL will not work for the CUDA runtime # noqa: F401 else: from tinygrad.runtime.cuda import CLBuffer, CLImage, CLProgram # type: ignore VALIDHACKS = getenv("VALIDHACKS", 0) # TODO: remove the need for this @@ -28,7 +28,7 @@ def split_float4(x): return sum([[Token(acc.tok+f".s{s}", Types.FLOAT) for s in range(4)] for acc in x], []) class CLASTKernel(ASTKernel): - code_for_op : Dict[Op, str] = { + code_for_op : Final[Dict[Op, str]] = { UnaryOps.NOOP: "(A)", UnaryOps.NEG: "(-(A))", UnaryOps.RELU: "max(A, (float)0.)", UnaryOps.GT0: "(A > 0.)" if CUDA else "((float)1.-step((float)0.,(-A)))", UnaryOps.EXP: "native_exp(A)" if NATIVE_EXPLOG else "exp(A)", @@ -38,7 +38,7 @@ class CLASTKernel(ASTKernel): BinaryOps.DIV: "(A/B)", BinaryOps.POW: "pow(A,B)", BinaryOps.CMPEQ: "(A==B)", ReduceOps.SUM: "A+=B", ReduceOps.MAX: "A=max(A,B)" } - start_for_op = {ReduceOps.SUM: "0.0", ReduceOps.MAX: "-INFINITY"} + start_for_op : Final[Dict[Op, str]] = {ReduceOps.SUM: "0.0", ReduceOps.MAX: "-INFINITY"} def image_idx(self, buf_index, idxy, validhacks=False): assert self.buftokens[buf_index].typ == Types.FLOAT4, f"image must be FLOAT4 {self.buftokens[buf_index]} {self.bufs[buf_index].st}" @@ -223,7 +223,7 @@ class CLASTKernel(ASTKernel): # STOP WASTING TIME WITH DOING THE RESHAPES AND PERMUTES BY HAND. KERNEL SEARCH IS THE ONLY WAY IT WILL EVER BE GOOD # group_for_reduce will have to be better first - def codegen(self): + def codegen(self) -> Callable: self.process() self.hand_coded_optimizations() @@ -257,12 +257,13 @@ class CLASTKernel(ASTKernel): # early ast accumulators : List[Token] = [Token("acc%d" % i, self.buftokens[0].typ) for i in range(self.buftokens[0].size())] - if self.reduceop: + if self.reduceop is not None: full_shape_candidates = [x.shape for x in self.sts if x.shape != self.sts[0].shape] full_shape : Tuple[int, ...] = self.sts[0].shape if len(full_shape_candidates) == 0 else full_shape_candidates[0] acc_offsets = self.buftokens[self.bufs.index(self.earlybufs[0])].acc_offsets() - self.kernel += [f"{accumulator.decltype()} {accumulator.tok} = {CLASTKernel.start_for_op[self.reduceop.op]};\n" for accumulator in accumulators] + assert self.reduceopop is not None + self.kernel += [f"{accumulator.decltype()} {accumulator.tok} = {CLASTKernel.start_for_op[self.reduceopop]};\n" for accumulator in accumulators] self.kernel += [f"for (int idx{i} = 0; idx{i} < {full_shape[i]}; idx{i}++) {{\n" for i in range(self.first_reduce+len(self.group_for_reduce), self.shape_len)] expanded_accumulators = split_float4(accumulators) if len(accumulators)*4 == len(acc_offsets) else accumulators self.kernel += [f"{x.tok};\n" for x in self.ast_parse(self.reduceop, [expanded_accumulators[off] for off in acc_offsets], do_reduce=True)] + ["}\n"] * (self.shape_len - (self.first_reduce + len(self.group_for_reduce))) @@ -283,14 +284,15 @@ class CLASTKernel(ASTKernel): self.reshape_and_permute(None, [i for i in range(self.shape_len) if i != self.first_reduce+1] + [self.first_reduce+1]) self.upcast() + assert self.reduceopop is not None self.kernel.append("if (mid_idx == 0) {\n") new_accumulators = [Token(f"output{i}", self.buftokens[0].typ) for i in range(len(accumulators))] for i,acc in enumerate(new_accumulators): self.kernel.append(f"{acc.decltype()} {acc.tok} = 0.0;\n") if self.upcast_in_mid_reduce: - self.kernel.append(f"for (int mid = 0; mid < {prod(self.group_for_reduce)//4}; mid++) {{ {CLASTKernel.code_for_op[self.reduceop.op].replace('A', acc.tok).replace('B', f'vload4(0, &temp{i}[mid*4])')}; }}\n") + self.kernel.append(f"for (int mid = 0; mid < {prod(self.group_for_reduce)//4}; mid++) {{ {CLASTKernel.code_for_op[self.reduceopop].replace('A', acc.tok).replace('B', f'vload4(0, &temp{i}[mid*4])')}; }}\n") else: - self.kernel.append(f"for (int mid = 0; mid < {prod(self.group_for_reduce)}; mid++) {{ {CLASTKernel.code_for_op[self.reduceop.op].replace('A', acc.tok).replace('B', f'temp{i}[mid]')}; }}\n") + self.kernel.append(f"for (int mid = 0; mid < {prod(self.group_for_reduce)}; mid++) {{ {CLASTKernel.code_for_op[self.reduceopop].replace('A', acc.tok).replace('B', f'temp{i}[mid]')}; }}\n") accumulators = new_accumulators # late ast @@ -326,17 +328,19 @@ class CLASTKernel(ASTKernel): class GPUBuffer(ExplicitExecAST): def __init__(self, shape:Union[ShapeTracker, Tuple[int, ...]], hostbuf:Optional[GPUBuffer]=None, backing:Optional[np.ndarray]=None, force_create=False): super().__init__(shape, hostbuf) - self._buf : Optional[CLBuffer] = hostbuf._buf if hostbuf is not None else None + self._buf : Optional[Union[CLImage, CLBuffer]] = hostbuf._buf if hostbuf is not None else None self._base_shape : Tuple[int, ...] = hostbuf._base_shape if hostbuf is not None else self.shape self._backing : Optional[np.ndarray] = hostbuf._backing if hostbuf is not None else backing # early copy in for large buffers if (self._backing is not None and self._backing.shape != (1,)) or force_create: self.cl + # TODO: refactor this to return self._buf and not import pyopencl @property - def cl(self): + def cl(self) -> pyopencl.Buffer: if self._buf is None: self._buf = CLImage(self._base_shape) if (len(self._base_shape) == 3 and self._base_shape[2] == 4 and IMAGE >= 2) else CLBuffer(4*prod(self._base_shape)) + assert self._buf is not None if self._backing is not None: self._buf.copyin(self._backing) self._backing = None diff --git a/tinygrad/ops.py b/tinygrad/ops.py index 193202ea6e..e91eafecec 100644 --- a/tinygrad/ops.py +++ b/tinygrad/ops.py @@ -1,7 +1,7 @@ from __future__ import annotations import numpy as np -from enum import Enum -from typing import Union, Type, NamedTuple, Tuple, Any, List +from enum import Enum, auto +from typing import Union, Type, NamedTuple, Tuple, Any, List, ClassVar import functools, operator from tinygrad.helpers import prod from tinygrad.shape import ShapeTracker @@ -10,12 +10,13 @@ from tinygrad.helpers import getenv DEBUG = getenv("DEBUG", 0) # these are the llops your accelerator must implement, along with toCpu -UnaryOps = Enum("UnaryOps", ["NOOP", "NEG", "RELU", "EXP", "LOG", "GT0", "RECIPROCAL"]) -BinaryOps = Enum("BinaryOps", ["ADD", "SUB", "MUL", "DIV", "POW", "CMPEQ"]) -ReduceOps = Enum("ReduceOps", ["SUM", "MAX"]) -MovementOps = Enum("MovementOps", ["RESHAPE", "PERMUTE", "EXPAND", "FLIP", "STRIDED", "PAD", "SHRINK"]) -ProcessingOps = Enum("ProcessingOps", ["CONV"]) -LoadOps = Enum("LoadOps", ["FROMCPU", "CONTIGUOUS"]) +# the Enum class doesn't work with mypy, this is static. sorry it's ugly +class UnaryOps(Enum): NOOP = auto(); NEG = auto(); RELU = auto(); EXP = auto(); LOG = auto(); GT0 = auto(); RECIPROCAL = auto() # noqa: E702 +class BinaryOps(Enum): ADD = auto(); SUB = auto(); MUL = auto(); DIV = auto(); POW = auto(); CMPEQ = auto() # noqa: E702 +class ReduceOps(Enum): SUM = auto(); MAX = auto() # noqa: E702 +class MovementOps(Enum): RESHAPE = auto(); PERMUTE = auto(); EXPAND = auto(); FLIP = auto(); STRIDED = auto(); PAD = auto(); SHRINK = auto() # noqa: E702 +class ProcessingOps(Enum): CONV = auto() # noqa: E702 +class LoadOps(Enum): FROMCPU = auto(); CONTIGUOUS = auto() # noqa: E702 Op = Union[UnaryOps, BinaryOps, ReduceOps, MovementOps, ProcessingOps, LoadOps] OpType = Union[Type[UnaryOps], Type[BinaryOps], Type[ReduceOps], Type[MovementOps], Type[ProcessingOps], Type[LoadOps]] @@ -30,6 +31,9 @@ class LazyOp(NamedTuple): # Any == Union[LazyBuffer, DeviceBuffer] def get_buffers(op:LazyOp) -> List[Any]: return functools.reduce(operator.add, [get_buffers(x) if isinstance(x, LazyOp) else [x] for x in op.src], []) def get_lazyops(op:LazyOp) -> List[LazyOp]: return functools.reduce(operator.add, [get_lazyops(x) for x in op.src if isinstance(x, LazyOp)], [op]) +def map_buffers(real_srcs, x:LazyOp) -> LazyOp: + if x in real_srcs: return map_buffers(real_srcs, real_srcs[x]) if isinstance(real_srcs[x], LazyOp) else real_srcs[x] + return LazyOp(x.op, tuple((map_buffers(real_srcs, y) if isinstance(y, LazyOp) else real_srcs[y]) for y in x.src), x.arg) # a placeholder class to extend by the exec classes class DeviceBuffer: @@ -63,7 +67,10 @@ class GenericExecAST(DeviceBuffer): # pylint: disable=abstract-method return ret class GlobalCounters: - global_ops, global_mem, time_sum = 0, 0, 0 + global_ops : ClassVar[int] = 0 + global_mem : ClassVar[int] = 0 + time_sum : ClassVar[int] = 0 + # TODO: reset method class GenericShape(GenericExecAST): # pylint: disable=abstract-method def __init__(self, shape, flops=0): self.shape, self.flops = shape, flops diff --git a/tinygrad/runtime/opencl.py b/tinygrad/runtime/opencl.py index 0a3317aa8b..87df92a62d 100644 --- a/tinygrad/runtime/opencl.py +++ b/tinygrad/runtime/opencl.py @@ -58,6 +58,9 @@ class CLImage: def __del__(self): CL.mem_used -= self.cl.row_pitch * self.cl.height + def copyin(self, b:np.ndarray): raise NotImplementedError("no copyin for CLImage") + def copyout(self, a:np.ndarray): raise NotImplementedError("no copyout for CLImage") + @functools.lru_cache(maxsize=None) class CLProgram: kernel_cnt : Dict[str, int] = defaultdict(int) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 74344cba8f..b2fb7ce958 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -3,13 +3,15 @@ from __future__ import annotations import inspect, functools, importlib, itertools import numpy as np from tinygrad.helpers import prod, argfix, make_pair -from typing import List, Tuple, Callable, Optional +from typing import List, Tuple, Callable, Optional, ClassVar, Type from tinygrad.lazy import Device, LazyBuffer # **** start with two base classes, Tensor and Function **** class Tensor: - training, no_grad = False, False + __deletable__ = ('_ctx',) + training : ClassVar[bool] = False + no_grad : ClassVar[bool] = False def __init__(self, data, device=Device.DEFAULT, requires_grad:Optional[bool]=None): if isinstance(data, list): @@ -40,14 +42,15 @@ class Tensor: return f"" @property - def shape(self): return self.lazydata.shape + def shape(self) -> Tuple[int, ...]: return self.lazydata.shape + # dtype handling was very broken. it's always float32 now @property - def dtype(self): return np.float32 + def dtype(self) -> type: return np.float32 @property - def device(self): return self.lazydata.device + def device(self) -> str: return self.lazydata.device # ***** data handlers **** @@ -63,11 +66,11 @@ class Tensor: return x def detach(self): return Tensor(self.lazydata, device=self.device, requires_grad=False) - def numpy(self): return np.array(self.lazydata.toCPU()) + def numpy(self) -> np.ndarray: return np.array(self.lazydata.toCPU()) # TODO: this keeps the legacy behavior working, remove it after refactor @property - def data(self): return self.numpy() + def data(self) -> np.ndarray: return self.numpy() # TODO: if things are realized this won't work def to_(self, device:str): @@ -168,11 +171,11 @@ class Tensor: val = [val] if not isinstance(val, (list, tuple)) else val assert sum(s is not None for s in val) <= len(self.shape) assert all(s.step is None or s.step == 1 for s in val if isinstance(s, slice)) - for i,(sz,s) in enumerate(zip(self.shape, (v for v in val if v is not None))): # Slicing only depends on ints + slices + for i,(sz,s) in enumerate(zip(self.shape, [v for v in val if v is not None])): # Slicing only depends on ints + slices if isinstance(s, int) and not (-sz <= s < sz): raise IndexError(f"index {s} is out of bounds for dimension {i} with size {sz}") new_slice.append((s%sz, s%sz+1) if isinstance(s, int) else (slcfix(s.start, sz, 0), slcfix(s.stop, sz, sz))) - for s,sz in zip(val, (self.shape[i-1] for i in itertools.accumulate(s is not None for s in val))): # Shape depends on slices + positions of Nones + for s,sz in zip(val, [self.shape[i-1] for i in itertools.accumulate([s is not None for s in val])]): # Shape depends on slices + positions of Nones if not isinstance(s, int): new_shape.append(1 if s is None else slcfix(s.stop, sz, sz) - slcfix(s.start, sz, 0)) new_shape += [self.shape[i] for i in range(len(new_slice), len(self.shape))] @@ -184,7 +187,7 @@ class Tensor: for y in args: assert len(y.shape) == len(self.shape) and all(y.shape[i] == s for i,s in enumerate(self.shape) if i != dim) catargs = [self] + list(args) - shape_cumsum = [0, *itertools.accumulate(y.shape[dim] for y in catargs)] + shape_cumsum = [0, *itertools.accumulate([y.shape[dim] for y in catargs])] slc = [[(0, s) for s in self.shape] for _ in catargs] for s,k in zip(slc, shape_cumsum): s[dim] = (-k, shape_cumsum[-1]-k) @@ -216,7 +219,7 @@ class Tensor: return cx.conv2d(cw, groups=groups).reshape(shape=out_shape_t).transpose(order=order) # TODO: what's the difference between dot and matmul? - dot = matmul + def dot(self:Tensor, w:Tensor): return self.matmul(w) # (padding_left, padding_right, padding_top, padding_bottom) def pad2d(self, padding:Tuple[int, ...]): return self.slice(arg = [(0,self.shape[0]), (0,self.shape[1]), (-padding[2],self.shape[2]+padding[3]), (-padding[0],self.shape[3]+padding[1])]) # type: ignore @@ -287,7 +290,7 @@ class Tensor: def sigmoid(self): return (1.0 + (-self).exp()).reciprocal() def elu(self, alpha=1.0): return self.relu() - alpha*(1-self.exp()).relu() def swish(self): return self * self.sigmoid() - silu = swish # The SiLU function is also known as the swish function. + def silu(self): return self.swish() # The SiLU function is also known as the swish function. def relu6(self): return self.relu() - (self-6).relu() def hardswish(self): return self * (self+3).relu6() * (1/6) def tanh(self): return 2.0 * ((2.0 * self).sigmoid()) - 1.0 @@ -307,12 +310,15 @@ class Tensor: shape_ret = tuple(max(sx, sy) for sx,sy in zip(x.shape, y.shape)) return fxn(x.expand(shape_ret), y.expand(shape_ret)) + @staticmethod + def _truediv(x, y): return x * (y.reciprocal() if isinstance(y, Tensor) else (1/y)) + # TODO: are these the only ones that can take number arguments? def add(self, x): return Tensor.broadcasted(Tensor._add, self, x) def sub(self, x): return Tensor.broadcasted(Tensor._sub, self, x) def mul(self, x): return Tensor.broadcasted(Tensor._mul, self, x) def pow(self, x): return Tensor.broadcasted(Tensor._pow, self, x) - def div(self, y): return self * (y.reciprocal() if isinstance(y, Tensor) else (1/y)) + def div(self, x): return Tensor._truediv(self, x) # ***** functional nn ops ***** @@ -349,30 +355,32 @@ class Function: # NOTE: it doesn't hurt to save this since the ctx will be freed fast without grad def save_for_backward(self, *x): self.saved_tensors.extend(x) - @classmethod - def apply(cls, *x:Tensor, **kwargs): - ctx = cls(x[0].device, *x) + @staticmethod + def apply(fxn:Type[Function], *x:Tensor, **kwargs): + ctx = fxn(x[0].device, *x) ret = Tensor(ctx.forward(*[t.lazydata for t in x], **kwargs), device=ctx.device, requires_grad=ctx.requires_grad) if ctx.requires_grad and not Tensor.no_grad: ret._ctx = ctx # used by autograd engine return ret # register functions to move between devices -for device in [device for device in Device.__dict__.keys() if device[0] != "_"]: - setattr(Tensor, f"{device.lower()}", functools.partialmethod(Tensor.to, Device.__dict__[device])) - setattr(Tensor, f"{device.lower()}_", functools.partialmethod(Tensor.to_, Device.__dict__[device])) +for device in [device for device in Device._buffers.keys() if device[0] != "_"]: + setattr(Tensor, f"{device.lower()}", functools.partialmethod(Tensor.to, device)) + setattr(Tensor, f"{device.lower()}_", functools.partialmethod(Tensor.to_, device)) # register all the mlops "math" operations -def register(name:str, fxn:Function): - setattr(Tensor, "_"+name if hasattr(Tensor, name) else name, lambda *args, **kwargs: fxn.apply(*args, **kwargs)) # doesn't work without lambda, pylint: disable=W0108 +def register(name:str, fxn:Type[Function]): + setattr(Tensor, "_"+name if hasattr(Tensor, name) else name, lambda *args, **kwargs: Function.apply(fxn, *args, **kwargs)) # doesn't work without lambda, pylint: disable=W0108 for name, cls in inspect.getmembers(importlib.import_module('tinygrad.mlops'), inspect.isclass): if name[0] != "_" and name != "Function" and not name.endswith("Ops"): register(name.lower(), cls) # register the operators -def register_op(name, fxn): +def register_op(name, fop): + if name in ['add', 'sub', 'mul', 'pow']: fxn = lambda x,y: Tensor.broadcasted(fop, x, y) + else: fxn = lambda x,y: fop(x,y) # pylint: disable=W0108 setattr(Tensor, f"__{name}__", fxn) - setattr(Tensor, f"__i{name}__", lambda self,x: self.assign(fxn(self,x))) - setattr(Tensor, f"__r{name}__", lambda self,x: fxn(x,self)) -for name in ['add', 'sub', 'mul', 'pow', 'matmul', 'truediv']: - register_op(name, getattr(Tensor, name if name != 'truediv' else 'div')) + setattr(Tensor, f"__i{name}__", lambda self,x: self.assign(fxn(self, x))) + setattr(Tensor, f"__r{name}__", lambda self,x: fxn(x, self)) +for name in ['add', 'sub', 'mul', 'pow', "truediv", "matmul"]: + register_op(name, getattr(Tensor, ("_"+name) if name != "matmul" else name))