diff --git a/README.md b/README.md index 20e62576de..c189faa6ec 100644 --- a/README.md +++ b/README.md @@ -158,12 +158,12 @@ You no longer need to write mlops for a new accelerator The autodiff stuff is all in mlops now so you can focus on the raw operations ``` -Buffer # class of memory on this device -unary_op (NOOP, NEG, NOT, EXP, LOG) # A -> A -reduce_op (SUM, MAX) # A -> B (smaller size, B has 1 in shape) -binary_op (ADD, SUB, MUL, DIV, POW, CMPEQ, MAX) # A + A -> A (all the same size) -movement_op (EXPAND, RESHAPE, PERMUTE, PAD, SHRINK, FLIP) # A -> B (different size) -fused_op [[optional]] (MULACC) # A * A -> B +Buffer # class of memory on this device +unary_op (NOOP, NEG, NOT, EXP, LOG) # A -> A +reduce_op (SUM, MAX) # A -> B (smaller size, B has 1 in shape) +binary_op (ADD, SUB, MUL, DIV, POW, CMPEQ, MAX) # A + A -> A (all the same size) +movement_op (EXPAND, RESHAPE, PERMUTE, PAD, SHRINK, STRIDE) # A -> B (different size) +fused_op [[optional]] (MULACC) # A * A -> B ``` ## ImageNet inference diff --git a/test/test_ops.py b/test/test_ops.py index ca941c927d..6d3078ebe2 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -269,6 +269,8 @@ class TestOps(unittest.TestCase): helper_test_op([(4,3,6,6)], lambda x: torch.flip(x, (0,1)), lambda x: x.flip(axis=(0,1))) helper_test_op([(4,3,6,6)], lambda x: torch.flip(x, (0,1,3)), lambda x: x.flip(axis=(0,1,3))) helper_test_op([(4,3,6,6)], lambda x: torch.flip(x, (3,)), lambda x: x.flip(axis=(3,))) + helper_test_op([(4,3,6,6)], lambda x: torch.flip(x, (0,1,3)).flip((0,)), lambda x: x.flip(axis=(0,1,3)).flip(0)) + helper_test_op([(4,3,6,6)], lambda x: torch.flip(x, (3,)), lambda x: x.flip(axis=(-1,))) def test_unsqueeze(self): helper_test_op([(4,3,6,6)], lambda x: torch.unsqueeze(x, 0), lambda x: x.unsqueeze(dim=0)) diff --git a/test/unit/test_shapetracker.py b/test/unit/test_shapetracker.py index 4a0a569ee4..35d063e9d8 100644 --- a/test/unit/test_shapetracker.py +++ b/test/unit/test_shapetracker.py @@ -23,31 +23,31 @@ class CheckingShapeTracker: def simplify(self): self.st.simplify() def reshape(self, new_shape): - self.st.reshape(new_shape) + self.st._reshape(new_shape) self.t = self.t.reshape(new_shape) def permute(self, axis): - self.st.permute(axis) + self.st._permute(axis) self.t = np.transpose(self.t, axis) def expand(self, new_shape): - self.st.expand(new_shape) + self.st._expand(new_shape) self.t = np.broadcast_to(self.t, new_shape) def flip(self, axis): - self.st.flip(axis) + self.st._stride(tuple(-1 if i in axis else 1 for i in range(len(self.shape)))) self.t = np.flip(self.t, axis) def shrink(self, arg): - self.st.shrink(arg) + self.st._shrink(arg) self.t = self.t[tuple([slice(x[0], x[1]) for x in arg])] def pad(self, arg): - self.st.pad(arg) + self.st._pad(arg) self.t = np.pad(self.t, arg, constant_values=-1) def stride(self, arg): - self.st.stride(arg) + self.st._stride(arg) self.t = self.t[tuple([slice(None, None, x) for x in arg])] def __getitem__(self, val): @@ -148,7 +148,7 @@ class TestSimplifyingShapeTracker(unittest.TestCase): class TestComplexShapeTracker(unittest.TestCase): def test_add_1s(self): - self.st = ShapeTracker((4, 4)) + self.st = CheckingShapeTracker((4, 4)) self.st.permute((1,0)) self.st.reshape((1,4,1,4,1)) assert not self.st.contiguous @@ -156,20 +156,20 @@ class TestComplexShapeTracker(unittest.TestCase): assert self.st.contiguous def test_permute_1s_simple(self): - self.st = ShapeTracker((1, 16, 9,9)) + self.st = CheckingShapeTracker((1, 16, 9,9)) self.st.permute((1,0,2,3)) assert self.st.contiguous - self.st = ShapeTracker((2, 16, 9,9)) + self.st = CheckingShapeTracker((2, 16, 9,9)) self.st.permute((1,0,2,3)) assert not self.st.contiguous def test_remove_1s_simple(self): - self.st = ShapeTracker((1, 16, 1, 1)) + self.st = CheckingShapeTracker((1, 16, 1, 1)) self.st.reshape((16,)) assert self.st.contiguous def test_remove_1s(self): - self.st = ShapeTracker((1, 4, 1, 4, 1)) + self.st = CheckingShapeTracker((1, 4, 1, 4, 1)) self.st.permute((0,3,2,1,4)) self.st.reshape((4,4)) assert not self.st.contiguous @@ -177,46 +177,46 @@ class TestComplexShapeTracker(unittest.TestCase): assert self.st.contiguous def test_permute_reshape(self): - self.st = ShapeTracker((4, 4)) + self.st = CheckingShapeTracker((4, 4)) self.st.permute((1,0)) self.st.reshape((2, 2, 2, 2)) # TODO: should also be tested by test_super_complex assert len(self.st.views) == 1 def test_factorize_split(self): - self.st = ShapeTracker((4, 4)) + self.st = CheckingShapeTracker((4, 4)) self.st.permute((1,0)) self.st.reshape((2, 2, 2, 2)) self.st.permute((2,3,0,1)) assert self.st.contiguous def test_factorize_combine(self): - self.st = ShapeTracker((4, 4, 4)) + self.st = CheckingShapeTracker((4, 4, 4)) self.st.permute((2, 0, 1)) self.st.reshape((4, 16)) self.st.permute((1, 0)) assert self.st.contiguous def test_factorize_combine_add_ones(self): - self.st = ShapeTracker((4, 4, 4)) + self.st = CheckingShapeTracker((4, 4, 4)) self.st.permute((2, 0, 1)) self.st.reshape((4, 16, 1, 1)) self.st.permute((1, 0, 2, 3)) assert self.st.contiguous def test_fancy_factorize(self): - self.st = ShapeTracker((32, 3, 3, 1)) + self.st = CheckingShapeTracker((32, 3, 3, 1)) self.st.reshape((8, 4, 3, 3)) assert len(self.st.views) == 1 def test_super_complex_2_fail(self): - self.st = ShapeTracker((4, 4, 4)) + self.st = CheckingShapeTracker((4, 4, 4)) self.st.permute((2, 0, 1)) self.st.reshape((16, 4)) assert len(self.st.views) != 1 def test_work(self): - self.st = ShapeTracker((64, 1024, 4)) + self.st = CheckingShapeTracker((64, 1024, 4)) self.st.reshape((1, 64, 128, 32)) self.st.permute((0, 3, 1, 2)) self.st.reshape((1, 32, 1, 64, 128)) @@ -224,7 +224,7 @@ class TestComplexShapeTracker(unittest.TestCase): assert self.st.contiguous def test_work2(self): - self.st = ShapeTracker((64, 1024, 4)) + self.st = CheckingShapeTracker((64, 1024, 4)) self.st.reshape((1, 64, 128, 32)) self.st.permute((0, 3, 1, 2)) self.st.reshape((1, 1, 32, 64, 128)) diff --git a/tinygrad/codegen/ast.py b/tinygrad/codegen/ast.py index 7acd47fe4f..0c964d9ec1 100644 --- a/tinygrad/codegen/ast.py +++ b/tinygrad/codegen/ast.py @@ -166,7 +166,7 @@ class ASTKernel: if mergeable: rets[j][-1] = (rets[j][-1][0] * shapes[j][i], strides[j][i]) else: rets[j].append((shapes[j][i], strides[j][i])) - for i,x in enumerate(rets): self.sts[i].reshape(tuple(y[0] for y in x)) + for i,x in enumerate(rets): self.sts[i]._reshape(tuple(y[0] for y in x)) self.first_reduce = get_first_reduce([x.shape for x in self.sts]) # this should be aware of the three parts to the shape @@ -175,8 +175,8 @@ class ASTKernel: # * the size outputted by each kernel def reshape_and_permute(self, new_shape_fxn, axis): for st in self.sts: - if new_shape_fxn is not None: st.reshape(tuple(new_shape_fxn(st.shape))) - if axis is not None: st.permute(tuple(axis)) + if new_shape_fxn is not None: st._reshape(tuple(new_shape_fxn(st.shape))) + if axis is not None: st._permute(tuple(axis)) # axis : the axis to pull from # amount : the amount to take diff --git a/tinygrad/lazy.py b/tinygrad/lazy.py index 1f00b41418..3411975c33 100644 --- a/tinygrad/lazy.py +++ b/tinygrad/lazy.py @@ -188,7 +188,7 @@ class LazyBuffer: if op in [MovementOps.RESHAPE, MovementOps.EXPAND, MovementOps.SHRINK]: return self.op.src[0].movement_op(op, arg) if op == MovementOps.PERMUTE: return self.op.src[0].movement_op(op, tuple(self.op.arg[i] for i in arg)) if op == MovementOps.PAD: return self.op.src[0].movement_op(op, tuple((b1+b2, e1+e2) for (b1,e1),(b2,e2) in zip(self.op.arg, arg))) - if op == MovementOps.FLIP: return self.op.src[0].movement_op(op, tuple(i for i in arg+self.op.arg if not (i in arg and i in self.op.arg))) + if op == MovementOps.STRIDE: return self.op.src[0].movement_op(op, tuple(i*j for i,j in zip(arg, self.op.arg))) # push permutes before reduce ops if op == MovementOps.PERMUTE and PUSH_PERMUTES and self.realized is None and self.optype == ReduceOps: diff --git a/tinygrad/mlops.py b/tinygrad/mlops.py index 109b638e94..aedcc43cf7 100644 --- a/tinygrad/mlops.py +++ b/tinygrad/mlops.py @@ -155,8 +155,8 @@ class Shrink(Function): class Flip(Function): def forward(self, x, axis): - self.axis = axis - return x.movement_op(MovementOps.FLIP, axis) + self.arg = tuple(-1 if i in axis else 1 for i in range(len(x.shape))) + return x.movement_op(MovementOps.STRIDE, self.arg) def backward(self, grad_output): - return grad_output.movement_op(MovementOps.FLIP, self.axis) + return grad_output.movement_op(MovementOps.STRIDE, self.arg) diff --git a/tinygrad/ops.py b/tinygrad/ops.py index 3ed4979543..bb5725d6e9 100644 --- a/tinygrad/ops.py +++ b/tinygrad/ops.py @@ -4,14 +4,13 @@ import numpy as np from enum import Enum, auto from typing import Union, Type, NamedTuple, Tuple, Any, List, ClassVar, Optional, Callable, Dict, TypeVar, Set from tinygrad.helpers import prod, DEBUG, getenv, DType, dtypes -from tinygrad.shape import ShapeTracker +from tinygrad.shape import ShapeTracker, MovementOps # these are the llops your accelerator must implement, along with toCpu # the Enum class doesn't work with mypy, this is static. sorry it's ugly class UnaryOps(Enum): NOOP = auto(); NEG = auto(); EXP = auto(); LOG = auto(); NOT = auto() # noqa: E702 class BinaryOps(Enum): ADD = auto(); SUB = auto(); MUL = auto(); DIV = auto(); POW = auto(); CMPEQ = auto(); MAX = auto() # noqa: E702 class ReduceOps(Enum): SUM = auto(); MAX = auto() # noqa: E702 -class MovementOps(Enum): RESHAPE = auto(); PERMUTE = auto(); EXPAND = auto(); FLIP = auto(); PAD = auto(); SHRINK = auto() # noqa: E702 class FusedOps(Enum): MULACC = auto() # noqa: E702 class LoadOps(Enum): FROMCPU = auto(); CONTIGUOUS = auto(); TOCPU = auto(); CUSTOM = auto() # noqa: E702 @@ -55,6 +54,11 @@ class RawBufferCopyIn(RawBuffer): ret.copyin(x) return ret +class RawBufferMapped(RawBufferCopyIn): + def _buffer(self) -> memoryview: raise NotImplementedError("must be implemented") + def toCPU(self) -> np.ndarray: return np.frombuffer(self._buffer(), dtype=self.dtype.np) + def copyin(self, x:np.ndarray) -> None: np.copyto(self.toCPU(), x.reshape(-1)) + class RawBufferCopyInOut(RawBufferCopyIn): def copyout(self, x:np.ndarray) -> None: raise NotImplementedError("must be implemented") @@ -130,7 +134,9 @@ class ASTRunner: if DEBUG >= 1: print(f"*** {GlobalCounters.kernel_count:4d} {self.name:20s} arg {len(rawbufs):3d} sz {str(self.global_size):18s} {str(self.local_size):12s} OPs {self.op_estimate/1e6:7.1f}M/{GlobalCounters.global_ops/1e9:7.2f}G mem {GlobalCounters.mem_used/1e9:5.2f} GB " + (str() if et is None else f"tm {et*1e6:9.2f}us/{GlobalCounters.time_sum_s*1e3:9.2f}ms ({self.op_estimate/(et*1e9):8.2f} GFLOPS, {self.mem_estimate/(et*1e9):6.2f} GB/s)")) - GlobalCounters.log_kernel(self.op_estimate, self.mem_estimate) + GlobalCounters.kernel_count += 1 + GlobalCounters.global_ops += self.op_estimate + GlobalCounters.global_mem += self.mem_estimate if getenv("EARLY_STOPPING") and GlobalCounters.kernel_count == getenv("EARLY_STOPPING"): exit(0) return et @@ -215,9 +221,4 @@ class GlobalCounters: mem_used : ClassVar[int] = 0 # NOTE: this is not reset cache : ClassVar[Optional[List[Tuple[Callable, Any]]]] = None @staticmethod - def reset(): GlobalCounters.global_ops, GlobalCounters.global_mem, GlobalCounters.time_sum_s, GlobalCounters.kernel_count, GlobalCounters.cache = 0,0,0.0,0,None - @staticmethod - def log_kernel(op_estimate:int, mem_estimate:int): - GlobalCounters.kernel_count += 1 - GlobalCounters.global_ops += op_estimate - GlobalCounters.global_mem += mem_estimate \ No newline at end of file + def reset(): GlobalCounters.global_ops, GlobalCounters.global_mem, GlobalCounters.time_sum_s, GlobalCounters.kernel_count, GlobalCounters.cache = 0,0,0.0,0,None \ No newline at end of file diff --git a/tinygrad/runtime/ops_clang.py b/tinygrad/runtime/ops_clang.py index 0fc6f287e2..e00ec6fb0f 100644 --- a/tinygrad/runtime/ops_clang.py +++ b/tinygrad/runtime/ops_clang.py @@ -1,18 +1,15 @@ import os, time, ctypes, hashlib, subprocess, platform -import numpy as np from collections import defaultdict from typing import Final, Dict from tinygrad.helpers import dtypes, DType -from tinygrad.ops import CompiledBuffer, RawBufferCopyIn +from tinygrad.ops import CompiledBuffer, RawBufferMapped from tinygrad.codegen.gpu import GPUCodegen, GPULanguage -class RawMallocBuffer(RawBufferCopyIn): +class RawMallocBuffer(RawBufferMapped): def __init__(self, size, dtype : DType): super().__init__(size, dtype) self._buf = ({dtypes.float32: ctypes.c_float, dtypes.float16: ctypes.c_int16}[dtype] * size)() - def _buffer(self): return self._buf - def copyin(self, x:np.ndarray): ctypes.memmove(self._buf, x.ctypes.data, x.size*np.dtype(x.dtype).itemsize) - def toCPU(self): return np.frombuffer(self._buf, dtype=self.dtype.np) + def _buffer(self): return memoryview(self._buf) class ClangProgram: kernel_cnt : Final[Dict[str, int]] = defaultdict(int) diff --git a/tinygrad/runtime/ops_cpu.py b/tinygrad/runtime/ops_cpu.py index 921b871408..748447d005 100644 --- a/tinygrad/runtime/ops_cpu.py +++ b/tinygrad/runtime/ops_cpu.py @@ -29,9 +29,9 @@ def einsum_mulacc(einsum, get_strides, expand): numpy_fxn_for_op : Dict[Op, Callable] = {**base_fxn_for_op, **{ UnaryOps.NOOP: np.ascontiguousarray, UnaryOps.EXP: np.exp, UnaryOps.LOG: np.log, BinaryOps.MAX: np.maximum, BinaryOps.CMPEQ: lambda x,y: (x==y).astype(np.float32), - MovementOps.FLIP: np.flip, MovementOps.PERMUTE: lambda x, order: x.transpose(order), - MovementOps.PAD: np.pad, MovementOps.EXPAND: np.broadcast_to, - FusedOps.MULACC: einsum_mulacc(lambda s,a,b: np.einsum(s, a.copy(), b.copy()), lambda x: x.strides, np.broadcast_to) + MovementOps.PERMUTE: lambda x, order: x.transpose(order), MovementOps.PAD: np.pad, MovementOps.EXPAND: np.broadcast_to, + MovementOps.STRIDE: lambda x, arg: x.__getitem__(tuple(slice(None, None, i) for i in arg)), + FusedOps.MULACC: einsum_mulacc(lambda s,a,b: np.einsum(s, a.copy(), b.copy()), lambda x: x.strides, np.broadcast_to), }} class CPUBuffer(InterpretedBuffer): diff --git a/tinygrad/runtime/ops_metal.py b/tinygrad/runtime/ops_metal.py index f8fa1563a4..65907bd69c 100644 --- a/tinygrad/runtime/ops_metal.py +++ b/tinygrad/runtime/ops_metal.py @@ -1,11 +1,10 @@ # pip3 install pyobjc-framework-Metal pyobjc-framework-Cocoa pyobjc-framework-libdispatch import os, subprocess, pathlib, functools import Metal, Cocoa, libdispatch # type: ignore -import numpy as np from typing import List, Any from tinygrad.codegen.gpu import GPUCodegen, GPULanguage from tinygrad.helpers import prod, getenv, DEBUG, DType -from tinygrad.ops import CompiledBuffer, RawBufferCopyIn +from tinygrad.ops import CompiledBuffer, RawBufferMapped METAL_XCODE = getenv("METAL_XCODE") @@ -19,20 +18,17 @@ class _METAL: return METAL.device.newCommandQueue() METAL = _METAL() -class RawMetalBuffer(RawBufferCopyIn): +class RawMetalBuffer(RawBufferMapped): def __init__(self, size:int, dtype:DType): super().__init__(size, dtype) self._cl = METAL.device.newBufferWithLength_options_(size*dtype.itemsize, Metal.MTLResourceStorageModeShared) def __del__(self): self._cl.release() super().__del__() - def _buffer(self): return self._cl.contents().as_buffer(self._cl.length()) - def _as_np(self): return np.frombuffer(self._buffer(), dtype=self.dtype.np) - def copyin(self, x:np.ndarray): np.copyto(self._as_np(), x.reshape(-1).data) - def toCPU(self) -> np.ndarray: + def _buffer(self): for cbuf in METAL.mtl_buffers_in_flight: cbuf.waitUntilCompleted() METAL.mtl_buffers_in_flight = [] - return self._as_np() # no copy! + return self._cl.contents().as_buffer(self._cl.length()) def unwrap(x): ret, err = x diff --git a/tinygrad/runtime/ops_torch.py b/tinygrad/runtime/ops_torch.py index fdc2dd0844..779b017798 100644 --- a/tinygrad/runtime/ops_torch.py +++ b/tinygrad/runtime/ops_torch.py @@ -8,7 +8,8 @@ torch_fxn_for_op : Dict[Op, Callable] = {**base_fxn_for_op, **{ UnaryOps.NOOP: lambda x: x.contiguous(), UnaryOps.EXP: lambda x: x.exp(), UnaryOps.LOG: lambda x: x.log(), BinaryOps.MAX: torch.maximum, BinaryOps.CMPEQ: lambda x,y: (x==y).float(), MovementOps.PAD: lambda x, padding: torch.nn.functional.pad(x, [item for sublist in padding[::-1] for item in sublist]), - FusedOps.MULACC: einsum_mulacc(lambda s,a,b: torch.einsum(s, a.float(), b.float()).type(a.dtype), lambda x: x.stride(), lambda x,s: x.expand(s)) + FusedOps.MULACC: einsum_mulacc(lambda s,a,b: torch.einsum(s, a.float(), b.float()).type(a.dtype), lambda x: x.stride(), lambda x,s: x.expand(s)), + MovementOps.STRIDE: lambda x, arg: x.__getitem__(tuple(slice(None, None, abs(i)) for i in arg)).flip([i for i,a in enumerate(arg) if a < 0]) }} device = torch.device("cuda:0" if torch.cuda.is_available() else ("mps" if getenv("MPS", 0) else "cpu")) diff --git a/tinygrad/shape/__init__.py b/tinygrad/shape/__init__.py index f4c78c2a11..772bb2cd22 100644 --- a/tinygrad/shape/__init__.py +++ b/tinygrad/shape/__init__.py @@ -1,10 +1,14 @@ # ShapeTracker allows movement operations to a buffer that don't require a copy to be made. from __future__ import annotations import functools -from typing import Tuple, Union, List, Optional, cast +from enum import Enum, auto +from typing import Tuple, Union, List, Optional, cast, Dict, Callable from tinygrad.helpers import prod, DEBUG from tinygrad.shape.symbolic import Variable, MulNode, NumNode, Node +# these ops live here +class MovementOps(Enum): RESHAPE = auto(); PERMUTE = auto(); EXPAND = auto(); PAD = auto(); SHRINK = auto(); STRIDE = auto() # noqa: E702 + @functools.lru_cache(maxsize=None) def to_shape_strides(shape:Tuple[int, ...], strides:Tuple[int, ...]) -> List[Tuple[int, int]]: assert len(shape) == len(strides) @@ -90,6 +94,8 @@ def merge_views(vm2:View, vm1:View) -> Optional[View]: return View(vm1.shape, tuple(new_strides), new_offset.b) if len(new_strides) == len(vm1.strides) else None class ShapeTracker: + dispatch : Dict[MovementOps, Callable] + def __init__(self, shape:Union[ShapeTracker, Tuple[int, ...]], views:Optional[List[ViewTypes]]=None): self.views : List[ViewTypes] = views if views is not None else (shape.views[:] if isinstance(shape, ShapeTracker) else [view_from_shape(shape)]) def __repr__(self): return f"ShapeTracker(shape={self.shape}, views={self.views})" @@ -133,13 +139,35 @@ class ShapeTracker: def expr_node(self, idx='idx', offset=0): return self._expr_idx(self.views[-1].expr_node(Variable(idx, 0, prod(self.shape)-1), offset)) - def movement_op(self, op, arg:Union[Tuple[int, ...], Tuple[Tuple[int, int], ...]]) -> ShapeTracker: - return getattr(self, str(op).split(".")[1].lower())(arg) def needs_valid(self) -> bool: return any(isinstance(v, ZeroView) for v in self.views) - def reshape(self, new_shape : Tuple[int, ...]) -> ShapeTracker: - assert isinstance(new_shape, tuple) + # *** under this line are the movement ops *** + + def __unsafe_resize(self, arg : Tuple[Tuple[int, int], ...]): + offset = sum([self.strides[i]*x for i,(x,_) in enumerate(arg)]) + self.views[-1] = View(tuple(y-x for x,y in arg), self.strides, self.offset+offset) + + def _pad(self, arg : Tuple[Tuple[int, int], ...]): + assert all((b>=0 and e>=0) for b,e in arg) and len(arg) == len(self.shape) + if all(b==0 and e==0 for b,e in arg): return self # ZeroView is expensive if we don't need it + zvarg = tuple((-b,s+e) for s,(b,e) in zip(self.shape, arg)) + zeroview = ZeroView(self.shape, zvarg) + self.__unsafe_resize(zvarg) + # if we add a ZeroView, we add another (stock) view also for modding + self.views += [zeroview, View(self.shape, strides_for_shape(self.shape))] + + def _shrink(self, arg : Tuple[Tuple[int, int], ...]): + assert all((b>=0 and e<=s) for s,(b,e) in zip(self.shape,arg)) and len(arg) == len(self.shape) + self.__unsafe_resize(arg) + + def _expand(self, new_shape : Tuple[int, ...]): + assert all(isinstance(x, int) for x in new_shape), f"non ints for expand in {new_shape}" + assert all(x == y or x == 1 for x,y in zip(self.shape, new_shape)), f"can't expand {self.shape} into {new_shape}" + strides : Tuple[int, ...] = tuple(s if x == y else 0 for s,(x,y) in zip(self.strides, zip(self.shape, new_shape))) + self.views[-1] = View(new_shape, strides, self.offset) + + def _reshape(self, new_shape : Tuple[int, ...]): if self.shape == new_shape: return self assert all(isinstance(x, int) and x != 0 for x in new_shape), f"shape must be ints and can't contain 0 {new_shape}" assert prod(self.shape) == prod(new_shape), f"can't reshape {self.shape} -> {new_shape}" @@ -150,61 +178,30 @@ class ShapeTracker: # NOTE: the last view in self.views is never a ZeroView if (merged_view := merge_views(cast(View, self.views[-1]), view)) is not None: self.views[-1] = merged_view else: self.views.append(view) - return self - def permute(self, axis : Tuple[int, ...]) -> ShapeTracker: - assert isinstance(axis, tuple) + def _permute(self, axis : Tuple[int, ...]): assert all(isinstance(x, int) and x >= 0 and x < len(self.shape) for x in axis), f"invalid permute {axis} for {self.shape}" assert len(set(axis)) == len(axis) and len(axis) == len(self.shape), f"can't permute {self.shape} with {axis}" self.views[-1] = View(tuple(self.shape[a] for a in axis), tuple(self.strides[a] for a in axis), self.offset) - return self - # TODO: this is a special case of slice with strides, remove it - # though it's nice that it can't change size - def flip(self, axis : Tuple[int, ...]) -> ShapeTracker: - return self.stride(tuple(-1 if i in axis else 1 for i in range(len((self.shape))))) - - # *** under this line are not invertible *** - - def _resize(self, arg : Tuple[Tuple[int, int], ...]): - offset = sum([self.strides[i]*x for i,(x,_) in enumerate(arg)]) - self.views[-1] = View(tuple(y-x for x,y in arg), self.strides, self.offset+offset) - - def pad(self, arg : Tuple[Tuple[int, int], ...]) -> ShapeTracker: - assert isinstance(arg, tuple) - assert all((b>=0 and e>=0) for b,e in arg) and len(arg) == len(self.shape) - if all(b==0 and e==0 for b,e in arg): return self # ZeroView is expensive if we don't need it - zvarg = tuple((-b,s+e) for s,(b,e) in zip(self.shape, arg)) - zeroview = ZeroView(self.shape, zvarg) - self._resize(zvarg) - # if we add a ZeroView, we add another (stock) view also for modding - self.views += [zeroview, View(self.shape, strides_for_shape(self.shape))] - return self - - def shrink(self, arg : Tuple[Tuple[int, int], ...]) -> ShapeTracker: - assert isinstance(arg, tuple) - assert all((b>=0 and e<=s) for s,(b,e) in zip(self.shape,arg)) and len(arg) == len(self.shape) - self._resize(arg) - return self - - def expand(self, new_shape : Tuple[int, ...]) -> ShapeTracker: - assert isinstance(new_shape, tuple) - assert all(isinstance(x, int) for x in new_shape), f"non ints for expand in {new_shape}" - assert all(x == y or x == 1 for x,y in zip(self.shape, new_shape)), f"can't expand {self.shape} into {new_shape}" - strides : Tuple[int, ...] = tuple(s if x == y else 0 for s,(x,y) in zip(self.strides, zip(self.shape, new_shape))) - self.views[-1] = View(new_shape, strides, self.offset) - return self - - # TODO: combine with flip? this is more generic than we need - def stride(self, mul : Tuple[int, ...]) -> ShapeTracker: - assert isinstance(mul, tuple) + # except for the negative case, you can build this from the others. invertible in the negative case + def _stride(self, mul : Tuple[int, ...]): assert all(isinstance(x, int) for x in mul) strides = tuple(z*m for z,m in zip(self.strides, mul)) new_shape = tuple((s+(abs(m)-1))//abs(m) for s,m in zip(self.shape, mul)) offset = sum([(s-1)*z for s,z,m in zip(self.shape, self.strides, mul) if m < 0]) self.views[-1] = View(new_shape, strides, self.offset + offset) + + # *** entry point for external *** + + def movement_op(self, op, arg:Union[Tuple[int, ...], Tuple[Tuple[int, int], ...]]) -> ShapeTracker: + assert isinstance(arg, tuple) and (len(arg) == len(self.shape) or op == MovementOps.RESHAPE), f"arg {arg} for {op} doesn't match dim of shape {self.shape}" + ShapeTracker.dispatch[op](self, arg) return self +# populate dispatch +ShapeTracker.dispatch = {op:getattr(ShapeTracker, "_"+str(op).split(".")[1].lower()) for op in MovementOps} + # returns the axes to create new_shape if new_shape can be created by combining axis from old_shape def get_contraction(old_shape:Tuple[int, ...], new_shape:Tuple[int, ...]): # Pre-allocate all groups. diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index b06b6f0e6b..b9bc953efd 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -191,7 +191,7 @@ class Tensor: return mlops.Reshape.apply(self, shape=tuple(-prod(self.shape) // prod(new_shape) if s == -1 else s for s in new_shape)) def expand(self, shape, *args) -> Tensor: return mlops.Expand.apply(self, shape=tuple(x if x != -1 else s for s,x in zip(self.shape, argfix(shape, *args)))) def permute(self, order, *args) -> Tensor: return mlops.Permute.apply(self, order=argfix(order, *args)) - def flip(self, axis, *args) -> Tensor: return mlops.Flip.apply(self, axis=argfix(axis, *args)) + def flip(self, axis, *args) -> Tensor: return mlops.Flip.apply(self, axis=[x if x >= 0 else x+len(self.shape) for x in argfix(axis, *args)]) def pad(self, arg:Tuple[Tuple[int, int], ...]) -> Tensor: return mlops.Pad.apply(self, arg=arg) if any(x != (0,0) for x in arg) else self def shrink(self, arg:Tuple[Tuple[int, int], ...]) -> Tensor: return mlops.Shrink.apply(self, arg=arg) if any(x != (0,s) for x,s in zip(arg, self.shape)) else self