diff --git a/test/test_ops.py b/test/test_ops.py index 1f7f1c3fca..b559762ce3 100644 --- a/test/test_ops.py +++ b/test/test_ops.py @@ -2,9 +2,8 @@ import torch import time import numpy as np import unittest -from tinygrad.tensor import Tensor, Device -from tinygrad.helpers import getenv -from tinygrad.lazy import IMAGE +from tinygrad.tensor import Tensor +from tinygrad.helpers import getenv, IMAGE FORWARD_ONLY = getenv("FORWARD_ONLY", 0) def helper_test_op(shps, torch_fxn, tinygrad_fxn=None, atol=1e-6, rtol=1e-3, grad_atol=1e-4, grad_rtol=1e-3, forward_only=False, vals=None, a=-0.5, b=3): diff --git a/test/test_shapetracker.py b/test/test_shapetracker.py index aeb7850e5f..002a93a694 100644 --- a/test/test_shapetracker.py +++ b/test/test_shapetracker.py @@ -178,7 +178,6 @@ class TestComplexShapeTracker(unittest.TestCase): def test_fancy_factorize(self): self.st = ShapeTracker((32, 3, 3, 1)) - self.st.strided(tuple(zip((32, 3, 3, 1), (1, 4096, 32, 1)))) self.st.reshape((8, 4, 3, 3)) assert len(self.st.views) == 1 diff --git a/tinygrad/graph.py b/tinygrad/graph.py index 6dace423f0..d9519cbc47 100644 --- a/tinygrad/graph.py +++ b/tinygrad/graph.py @@ -7,7 +7,7 @@ except ImportError: nx = None # graph won't work from collections import defaultdict from typing import Dict, List, Optional -from tinygrad.ops import DeviceBuffer, UnaryOps, BinaryOps, ReduceOps, MovementOps, ProcessingOps, LoadOps, FusedOps, Op, OpType, LazyOp, get_buffers, get_lazyops +from tinygrad.ops import DeviceBuffer, UnaryOps, BinaryOps, ReduceOps, MovementOps, LoadOps, FusedOps, Op, OpType, LazyOp, get_buffers, get_lazyops from tinygrad.helpers import getenv, DEBUG GRAPH, PRUNEGRAPH, GRAPHPATH = getenv("GRAPH", 0), getenv("PRUNEGRAPH", 0), getenv("GRAPHPATH", "/tmp/net") @@ -49,13 +49,13 @@ def log_op(ret : DeviceBuffer, ast : LazyOp, show_graph : Optional[bool] = None) if len(inp) == 1 and inp[0] == ret: if nm(ret) in G.nodes: G.nodes[nm(ret)]['style'] += ', bold' return # don't log self loops - oporder = [LoadOps, FusedOps, ProcessingOps, ReduceOps, BinaryOps, UnaryOps, MovementOps] + oporder = [LoadOps, FusedOps, ReduceOps, BinaryOps, UnaryOps, MovementOps] optype = type(sorted(op, key=lambda x: oporder.index(type(x)))[0]) cnts[optype] += 1 if DEBUG >= 3: print(f"{op} : {', '.join([f'{x.shape}-<{nm(x)}>' for x in inp])} -> {ret.shape}-<{nm(ret)}>") if show_graph: - top_colors = {LoadOps: '#FFFF80', UnaryOps: "#c0c0c0", ReduceOps: "#8080ff", BinaryOps: "#c0c0c0", MovementOps: "#80ff80", ProcessingOps: "#ff8080", FusedOps: "#ff8080"} + top_colors = {LoadOps: '#FFFF80', UnaryOps: "#c0c0c0", ReduceOps: "#8080ff", BinaryOps: "#c0c0c0", MovementOps: "#80ff80", FusedOps: "#ff8080"} dashed = (optype == LoadOps and hasattr(ret, "_backing")) or (hasattr(ret, "st") and not ret.st.contiguous) # type: ignore for x in inp: diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index aa02b8d2bb..22eb655a87 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -1,4 +1,3 @@ -from collections import namedtuple import os, math, functools, time from typing import Tuple, Union, List @@ -20,31 +19,8 @@ class Timing(object): @functools.lru_cache(maxsize=None) def getenv(key, default=0): return type(default)(os.getenv(key, default)) DEBUG = getenv("DEBUG", 0) +IMAGE = getenv("IMAGE", 0) def shape_to_axis(old_shape, new_shape): assert len(old_shape) == len(new_shape), "reduce shapes must have same dimensions" return tuple([i for i,(a,b) in enumerate(zip(old_shape, new_shape)) if a != b]) - -ConvArgs = namedtuple('ConvArgs', ['H', 'W', 'groups', 'rcout', 'cin', 'oy', 'ox', 'iy', 'ix', 'sy', 'sx', 'bs', 'cout', 'py', 'py_', 'px', 'px_', 'dy', 'dx', 'out_shape']) -def get_conv_args(x_shape, w_shape, stride=1, groups=1, padding=0, dilation=1, out_shape=None): - # TODO: https://docs.nvidia.com/deeplearning/performance/dl-performance-convolutional/index.html#tensor-layout - cout,cin,H,W = w_shape - sy,sx = make_pair(stride) - px,px_,py,py_ = [padding]*4 if isinstance(padding, int) else (padding if len(padding) == 4 else [padding[1], padding[1], padding[0], padding[0]]) - dy,dx = make_pair(dilation) - bs,cin_,iy,ix = x_shape - - # this can change px_ and py_ to make the out_shape right - # TODO: copy padding names from http://nvdla.org/hw/v1/ias/unit_description.html - if out_shape is not None: - py_ = (out_shape[2] - 1) * sy + 1 + dy * (H-1) - iy - py - px_ = (out_shape[3] - 1) * sx + 1 + dx * (W-1) - ix - px - - # TODO: should be easy to support asymmetric padding by changing output size - # https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html describes these sizes well - oy = (iy + py + py_ - dy * (H-1) - 1)//sy + 1 - ox = (ix + px + px_ - dx * (W-1) - 1)//sx + 1 - if cin*groups != cin_: - raise TypeError(f"Input Tensor shape {x_shape} does not match the shape of the weights {w_shape}. ({cin*groups} vs. {cin_})") - assert cout % groups == 0 and (out_shape is None or out_shape == (bs, cout, oy, ox)) - return ConvArgs(H, W, groups, cout//groups, cin, oy, ox, iy, ix, sy, sx, bs, cout, py, py_, px, px_, dy, dx, (bs, cout, oy, ox)) diff --git a/tinygrad/image.py b/tinygrad/image.py index 1fd05c08ee..e853937210 100644 --- a/tinygrad/image.py +++ b/tinygrad/image.py @@ -1,10 +1,7 @@ -from tinygrad.tensor import HLOP -from tinygrad.helpers import getenv - -IMAGE = getenv("IMAGE", 0) +from tinygrad.helpers import IMAGE def image_conv2d_decorator(normal_conv): - if not HLOP or IMAGE == 0: return normal_conv + if IMAGE == 0: return normal_conv def image_conv2d(self, weight, bias=None, groups=1, stride=1, dilation=1, padding=0): (bs,_,iy,ix), (cout,cin,H,W) = self.shape, weight.shape diff --git a/tinygrad/lazy.py b/tinygrad/lazy.py index 9a6dcb75d5..a5fb4335f7 100644 --- a/tinygrad/lazy.py +++ b/tinygrad/lazy.py @@ -2,9 +2,9 @@ from __future__ import annotations from typing import Optional, Tuple, Union, List, Dict, Any, ClassVar, Type import sys, weakref, importlib, inspect from weakref import WeakValueDictionary -from tinygrad.helpers import ConvArgs, prod, DEBUG +from tinygrad.helpers import prod, DEBUG from tinygrad.shape import ShapeTracker -from tinygrad.ops import DeviceBuffer, UnaryOps, BinaryOps, ReduceOps, MovementOps, ProcessingOps, LoadOps, OpType, LazyOp, get_buffers, get_lazyops, map_buffers, GenericExecAST +from tinygrad.ops import DeviceBuffer, UnaryOps, BinaryOps, ReduceOps, MovementOps, LoadOps, OpType, LazyOp, get_buffers, get_lazyops, map_buffers from tinygrad.graph import log_op from tinygrad.helpers import getenv @@ -12,8 +12,6 @@ from tinygrad.helpers import getenv sys.setrecursionlimit(10000) OPT = getenv("OPT", 2) -NOCONV = getenv("NOCONV", 0) -IMAGE = getenv("IMAGE", 0) LAZY = getenv("LAZY", 1) def get_buffer(name, base='tinygrad.llops'): @@ -51,18 +49,16 @@ def _ast_binaryops(self:LazyBuffer) -> LazyOp: real_srcs : Dict[LazyBuffer, Union[None, LazyOp, LazyBuffer]] = {x:None for x in get_buffers(self.op)} if DEBUG >= 3: for k,x in zip(real_srcs.keys(), map(get_movementroot_contiguous, real_srcs.keys())): - if x.optype in [ProcessingOps,ReduceOps] and x.realized is None: - print("\nHIT", k,x, "UNFOLDABLE" if len(k.children) > 1 or len(x.children) > 1 else "") + if x.optype == ReduceOps and x.realized is None: + print("\nHIT", k,x, "UNFOLDABLE" if len(k.children) > 1 or len(x.children) > 1 else str()) for tk in k.children: print("k", tk) for tx in x.children: print("x", tx) # NOTE: contiguous does not always mean the same size with SHRINK. this is still mergeable but requires more thought how - psrcs : List[Tuple[LazyBuffer, LazyBuffer]] = [(k,x) for k,x in zip(real_srcs.keys(), map(get_movementroot_contiguous, real_srcs.keys())) if x.optype in [ProcessingOps,ReduceOps] and x.realized is None and prod(k.shape) == prod(x.shape) and len(x.children) <= 1 and len(k.children) <= 1] + psrcs : List[Tuple[LazyBuffer, LazyBuffer]] = [(k,x) for k,x in zip(real_srcs.keys(), map(get_movementroot_contiguous, real_srcs.keys())) if x.optype == ReduceOps and x.realized is None and prod(k.shape) == prod(x.shape) and len(x.children) <= 1 and len(k.children) <= 1] intermediate_shape : Tuple[int, ...] = self.shape if len(psrcs) == 1 and MERGE_ONE_REDUCE_INTO_ELEMENTWISE: if DEBUG >= 3: print("FOLDING", psrcs[0]) - if psrcs[0][1].optype == ProcessingOps: - top = psrcs[0][1].op # _ast_processingops - elif psrcs[0][1].optype == ReduceOps: + if psrcs[0][1].optype == ReduceOps: top = _ast_reduceops(psrcs[0][1]) real_srcs[psrcs[0][0]] = top real_srcs.update({x:x for x in get_buffers(top)}) # the reduce op buffers are not modified @@ -151,7 +147,6 @@ class LazyBuffer: real_src = src.realize(self.device) self.realized = real_src.movement_op(self.op.op, self.op.arg) ast = LazyOp(self.op.op, (real_src, )) - elif self.optype == ProcessingOps: ast = self.op # no ast modifications for ProcessingOps elif self.optype == ReduceOps: ast = _ast_reduceops(self) elif self.optype == BinaryOps: ast = _ast_binaryops(self) @@ -200,11 +195,10 @@ class LazyBuffer: local_st = ShapeTracker(self.shape).movement_op(op, arg) # instant nops - if local_st.contiguous and self.shape == local_st.shape and op != MovementOps.STRIDED: - return self + if local_st.contiguous and self.shape == local_st.shape: return self # two ops in a row is one op. merge them if unresolved - if self.realized is None and self.op.op == op and op != MovementOps.STRIDED: + if self.realized is None and self.op.op == op: # TODO: why is deleting self from children needed? shouldn't GC do it? self.op.src[0].children.discard(self) if op in [MovementOps.RESHAPE, MovementOps.EXPAND, MovementOps.SHRINK]: @@ -264,13 +258,8 @@ class LazyBuffer: return self.op.src[0].movement_op(MovementOps.PERMUTE, tuple(new_arg)) \ .movement_op(MovementOps.RESHAPE, ShapeTracker(self.st).movement_op(op, arg).shape) - # some strideds are actually just reshapes - # NOTE: due to how strided works, we have to check the parent to be contiguous also - if op == MovementOps.STRIDED and local_st.contiguous and self.st.contiguous: - return self.movement_op(MovementOps.RESHAPE, tuple(i for i,_ in arg)) - # if this MovementOp is being applied to a BinaryOp, apply the MovementOp to all the BinaryOp inputs instead. NOTE: UnaryOps is never an OpType - if SHUFFLE_MOVEMENT_OPS and self.optype == BinaryOps and self.realized is None and len(self.children) == 0 and op not in [MovementOps.EXPAND, MovementOps.STRIDED] and (op != MovementOps.PAD or all(x.op != BinaryOps.DIV for x in get_lazyops(self.op))): + if SHUFFLE_MOVEMENT_OPS and self.optype == BinaryOps and self.realized is None and len(self.children) == 0 and op != MovementOps.EXPAND and (op != MovementOps.PAD or all(x.op != BinaryOps.DIV for x in get_lazyops(self.op))): return replace_with_movement_op(self.op, op, arg) # create the buffer @@ -286,100 +275,6 @@ class LazyBuffer: return ret - def processing_op(self:LazyBuffer, op:ProcessingOps, w:LazyBuffer, C:ConvArgs) -> LazyBuffer: - x = self - - if IMAGE >= 1: - w = w.movement_op(MovementOps.RESHAPE, (C.groups, C.rcout, C.cin, C.H, C.W)) - # TODO: moving the x reshape here creates more views? - - # hack for non multiples of 4 on C.cin - if C.cin % 4 != 0 and not (C.cin == 1 and C.groups%4 == 0): - added_input_channels = 4 - (C.cin % 4) - w = w.movement_op(MovementOps.PAD, tuple((0, added_input_channels) if i == 2 else (0, 0) for i in range(len(w.shape)))) - x = x.movement_op(MovementOps.RESHAPE, (C.bs, C.groups, C.cin, C.iy, C.ix)) - x = x.movement_op(MovementOps.PAD, tuple((0, added_input_channels) if i == 2 else (0, 0) for i in range(len(x.shape)))) - C = C._replace(cin = C.cin + added_input_channels) - x = x.movement_op(MovementOps.RESHAPE, (C.bs, C.groups*C.cin, C.iy, C.ix)) - - # hack for non multiples of 4 on C.rcout - added_output_channels = 0 - if C.rcout % 4 != 0 and not (C.rcout == 1 and C.groups%4 == 0): - added_output_channels = 4 - (C.rcout % 4) - w = w.movement_op(MovementOps.PAD, tuple((0, added_output_channels) if i == 1 else (0, 0) for i in range(len(w.shape)))) - C = C._replace(rcout = C.rcout + added_output_channels, cout = C.groups * (C.rcout + added_output_channels)) - - # packed - x = x.movement_op(MovementOps.PERMUTE, (0,2,3,1)) - x = x.movement_op(MovementOps.RESHAPE, (C.bs*C.iy, C.ix*C.groups*C.cin//4, 4)) - - if C.cin == 1: # depthwise - w = w.movement_op(MovementOps.RESHAPE, (C.cout//4,4,C.H*C.W)) - w = w.movement_op(MovementOps.PERMUTE, (0,2,1)) - else: - w = w.movement_op(MovementOps.RESHAPE, (C.cout//4,4,C.cin//4,4,C.H,C.W)) - w = w.movement_op(MovementOps.PERMUTE, (0,4,2,5,1,3)) - w = w.movement_op(MovementOps.RESHAPE, (C.cout//4, C.H * C.cin//4 * C.W * 4, 4)) - - # contiguous creates the image, and early realize static weights - x, w = x.contiguous(), w.contiguous() - if get_single_root(w).realized: w.realize() - - # set up the conv from (C.bs*C.iy, C.ix*C.groups*C.cin//4, 4), pad, stride, and expand - x = x.movement_op(MovementOps.RESHAPE, (C.bs, C.iy, C.ix, C.groups, C.cin)) - x = x.slice(((0, x.shape[0]), (-C.py, x.shape[1]+C.py_), (-C.px, x.shape[2]+C.px_), (0, x.shape[3]), (0, x.shape[4]))) - x = x.movement_op(MovementOps.STRIDED, ( - (C.bs, x.shape[1]*x.shape[2]*C.groups*C.cin), - (C.oy, C.sy*x.shape[2]*C.groups*C.cin), (C.ox, C.sx*C.groups*C.cin), - (C.groups, C.cin), (1, 1), (1, 1), - (C.H, C.dy*x.shape[2]*C.groups*C.cin), (C.W, C.dx*C.groups*C.cin), (C.cin//4 if C.cin >= 4 else 1, 4), (4 if C.cin >= 4 else 1, 1) - )) - x = x.movement_op(MovementOps.EXPAND, (C.bs, C.oy, C.ox, C.groups, C.rcout//4 if C.rcout >= 4 else 1, 4 if C.rcout >= 4 else 1, C.H, C.W, x.shape[-2], x.shape[-1])) - x = x.movement_op(MovementOps.RESHAPE, (C.bs, C.oy, C.ox, C.cout//4, 4, C.H, C.W, x.shape[-2], x.shape[-1])) - - # set up the weights from (C.cout//4, C.H * C.cin//4 * C.W * 4, 4) - w = w.movement_op(MovementOps.RESHAPE, (C.cout//4, C.H, C.cin//4 if C.cin >= 4 else 1, C.W, 4, 4 if C.cin >= 4 else 1)) - w = w.movement_op(MovementOps.PERMUTE, (0,4,1,3,2,5)) - w = w.movement_op(MovementOps.RESHAPE, (1, 1, 1, C.cout//4, 4, C.H, C.W, w.shape[-2], w.shape[-1])) - w = w.movement_op(MovementOps.EXPAND, (C.bs, C.oy, C.ox, C.cout//4, 4, C.H, C.W, w.shape[-2], w.shape[-1])) - - # now do the conv in this space and force it to be an image - ret = x.binary_op(BinaryOps.MUL, w).reduce_op(ReduceOps.SUM, (C.bs, C.oy, C.ox, C.cout//4, 4, 1, 1, 1, 1)) - ret = ret.movement_op(MovementOps.RESHAPE, (C.bs*C.oy, C.ox*C.cout//4, 4)) - # NOTE: right now, this can't always be an image, because the tests fail when we try to access it if it was SHRINKed - if IMAGE >= 3: ret = ret.contiguous() # but if IMAGE >= 3, you can do anything - - # undo hack for non multiples of 4 on C.rcout - if added_output_channels != 0: - ret = ret.movement_op(MovementOps.RESHAPE, (C.bs, C.oy, C.ox, C.groups, C.rcout)) - ret = ret.movement_op(MovementOps.SHRINK, tuple((0, s-added_output_channels) if i == 4 else (0, s) for i,s in enumerate(ret.shape))) - C = C._replace(rcout = C.rcout - added_output_channels, cout = C.groups * (C.rcout - added_output_channels)) - - ret = ret.movement_op(MovementOps.RESHAPE, (C.bs, C.oy, C.ox, C.cout)) - ret = ret.movement_op(MovementOps.PERMUTE, (0,3,1,2)) - return ret - - # add padding if the backend can't handle it - if NOCONV or (not getattr(x.dbuffer, "SUPPORTS_PADDING", False) and not (getattr(x.dbuffer, "SUPPORTS_SIMPLE_PADDING", False) and C.px == C.px_ and C.py == C.py_ and C.px >= 0 and C.py >= 0)): - x = x.slice(((0, x.shape[0]), (0, x.shape[1]), (-C.py, x.shape[2]+C.py_), (-C.px, x.shape[3]+C.px_))) - C = C._replace(px=0, px_=0, py=0, py_=0) - - if NOCONV or not issubclass(x.dbuffer, GenericExecAST): - # universal conv, just mul and reduce - x = x.movement_op(MovementOps.STRIDED, ( - (C.bs, C.groups*C.cin*x.shape[2]*x.shape[3]), (C.groups, C.cin*x.shape[2]*x.shape[3]), - (1, 1), (C.oy, C.sy*x.shape[3]), (C.ox, C.sx), - (C.cin, x.shape[2]*x.shape[3]), (C.H, C.dy*x.shape[3]), (C.W, C.dx))) - #if C.H <= 3 and C.W <= 3: # max 9x the RAM overhead, this is im2col - # x = x.contiguous() - x = x.movement_op(MovementOps.EXPAND, (C.bs, C.groups, C.rcout, C.oy, C.ox, C.cin, C.H, C.W)) - w = w.movement_op(MovementOps.RESHAPE, (1, C.groups, C.rcout, 1, 1, C.cin, C.H, C.W)) \ - .movement_op(MovementOps.EXPAND, (C.bs, C.groups, C.rcout, C.oy, C.ox, C.cin, C.H, C.W)) - return x.binary_op(BinaryOps.MUL, w).reduce_op(ReduceOps.SUM, (C.bs, C.groups, C.rcout, C.oy, C.ox, 1, 1, 1)) \ - .movement_op(MovementOps.RESHAPE, (C.bs, C.cout, C.oy, C.ox)) - else: - return LazyBuffer(x.device, C.out_shape, ProcessingOps, LazyOp(op, (x, w), C)) - def elementwise_op(op:Union[UnaryOps, BinaryOps], *srcs:LazyBuffer) -> LazyBuffer: out_device, out_shape = srcs[0].device, srcs[0].shape diff --git a/tinygrad/llops/ops_cpu.py b/tinygrad/llops/ops_cpu.py index 4216b1d50a..e5ca55617b 100644 --- a/tinygrad/llops/ops_cpu.py +++ b/tinygrad/llops/ops_cpu.py @@ -1,7 +1,7 @@ import numpy as np import operator from typing import ClassVar, Callable, Dict -from tinygrad.ops import UnaryOps, BinaryOps, MovementOps, ReduceOps, ProcessingOps, FusedOps, GenericExecAST, Op +from tinygrad.ops import UnaryOps, BinaryOps, MovementOps, ReduceOps, FusedOps, GenericExecAST, Op from tinygrad.helpers import shape_to_axis base_fxn_for_op : Dict[Op, Callable] = { @@ -12,16 +12,6 @@ base_fxn_for_op : Dict[Op, Callable] = { MovementOps.SHRINK: lambda x, arg: x[tuple(slice(p[0], p[1], None) for p in arg)], } -def numpy_strided(x,arg): return np.lib.stride_tricks.as_strided(x.ravel().reshape(x.shape), shape=[y[0] for y in arg], strides=[y[1]*x.dtype.itemsize for y in arg]) -def numpy_conv(x,w,C): - assert C.px == 0 and C.px_ == 0 and C.py == 0 and C.py_ == 0, "padding in conv is not supported" - tx = numpy_strided(x, ( - (C.bs, C.groups*C.cin*x.shape[2]*x.shape[3]), (C.groups, C.cin*x.shape[2]*x.shape[3]), - (C.oy, C.sy*x.shape[3]), (C.ox, C.sx), (C.cin, x.shape[2]*x.shape[3]), (C.H, C.dy*x.shape[3]), (C.W, C.dx))) - tw = w.reshape(C.groups, C.rcout, C.cin, C.H, C.W) - out = np.einsum("nGhwCHW, GkCHW -> nGkhw", tx.ravel().reshape(tx.shape), tw.ravel().reshape(tw.shape)) - return out.reshape(C.bs, C.groups*C.rcout, C.oy, C.ox) - def einsum_mulacc(einsum, get_strides, expand): def einscripts(x): return ''.join(["abcdefghijklmnopqrstuvwxyz"[i] for i in x]) def axes_slice(strides): return [i for i in range(len(strides)) if strides[i] != 0], tuple(slice(None) if strides[i] != 0 else 0 for i in range(len(strides))) @@ -37,8 +27,7 @@ numpy_fxn_for_op : Dict[Op, Callable] = {**base_fxn_for_op, **{ BinaryOps.MAX: np.maximum, BinaryOps.CMPEQ: lambda x,y: (x==y).astype(np.float32), MovementOps.FLIP: lambda x, axis: np.flip(x, axis), MovementOps.PERMUTE: lambda x, order: x.transpose(order), MovementOps.PAD: lambda x, padding: np.pad(x, padding), MovementOps.EXPAND: lambda x, new_shape: np.broadcast_to(x, new_shape), - MovementOps.STRIDED: numpy_strided, ProcessingOps.CONV: numpy_conv, - FusedOps.MULACC: einsum_mulacc(lambda s,a,b: np.einsum(s, a.copy(), b.copy()), lambda x: x.strides, lambda x,s: np.broadcast_to(x,s)) + FusedOps.MULACC: einsum_mulacc(lambda s,a,b: np.einsum(s, a.copy(), b.copy()), lambda x: x.strides, np.broadcast_to) }} class CPUBuffer(GenericExecAST): diff --git a/tinygrad/llops/ops_gpu.py b/tinygrad/llops/ops_gpu.py index 7b34b04530..39b6241628 100644 --- a/tinygrad/llops/ops_gpu.py +++ b/tinygrad/llops/ops_gpu.py @@ -2,10 +2,9 @@ from __future__ import annotations import numpy as np import math from typing import List, Tuple, Optional, Dict, Union, Set, Final, Callable -from tinygrad.helpers import prod, DEBUG +from tinygrad.helpers import prod, DEBUG, IMAGE from tinygrad.ops import UnaryOps, BinaryOps, ReduceOps, MovementOps, LazyOp, Op, ExplicitExecAST, GlobalCounters from tinygrad.ast import ASTKernel, Token, Types -from tinygrad.lazy import IMAGE from tinygrad.shape import ShapeTracker from tinygrad.shape.symbolic import Node, ModNode, DivNode, render_python # div is different in cl than python diff --git a/tinygrad/llops/ops_torch.py b/tinygrad/llops/ops_torch.py index fb95280b5f..f881fc37cd 100644 --- a/tinygrad/llops/ops_torch.py +++ b/tinygrad/llops/ops_torch.py @@ -1,6 +1,6 @@ import torch from typing import ClassVar, Final, Dict, Callable -from tinygrad.ops import UnaryOps, BinaryOps, MovementOps, ProcessingOps, FusedOps, GenericExecAST, Op +from tinygrad.ops import UnaryOps, BinaryOps, MovementOps, FusedOps, GenericExecAST, Op from tinygrad.helpers import getenv from tinygrad.llops.ops_cpu import base_fxn_for_op, einsum_mulacc @@ -8,8 +8,6 @@ torch_fxn_for_op : Dict[Op, Callable] = {**base_fxn_for_op, **{ 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]), - MovementOps.STRIDED: lambda x, arg: x.contiguous().as_strided([y[0] for y in arg], [y[1] for y in arg]), - ProcessingOps.CONV: lambda x,w,C: C.px == C.px_ and C.py == C.py_ and torch.conv2d(x, w, stride=(C.sy, C.sx), groups=C.groups, dilation=(C.dy, C.dx), padding=(C.py, C.px)), FusedOps.MULACC: einsum_mulacc(torch.einsum, lambda x: x.stride(), lambda x,s: x.expand(s)) }} diff --git a/tinygrad/mlops.py b/tinygrad/mlops.py index 995696587b..19c5f50bb7 100644 --- a/tinygrad/mlops.py +++ b/tinygrad/mlops.py @@ -1,5 +1,5 @@ -from tinygrad.helpers import prod, argsort, get_conv_args -from tinygrad.ops import UnaryOps, BinaryOps, ReduceOps, MovementOps, ProcessingOps +from tinygrad.helpers import prod, argsort +from tinygrad.ops import UnaryOps, BinaryOps, ReduceOps, MovementOps from tinygrad.tensor import Function class Contiguous(Function): @@ -157,37 +157,3 @@ class Flip(Function): def backward(self, grad_output): return grad_output.movement_op(MovementOps.FLIP, self.axis) - -# ************* processing ops ************* - -class Conv2D(Function): - def forward(self, x, w, stride=1, groups=1, dilation=1, padding=0): - self.C = get_conv_args(x.shape, w.shape, stride, groups, dilation=dilation, padding=padding) - self.save_for_backward(x,w) - return x.processing_op(ProcessingOps.CONV, w, self.C) - - def backward(self, grad_output): - x, w = self.saved_tensors - C = self.C # conv args from the context - dx, dw = None, None - - if self.needs_input_grad[0]: # compute derivative of inputs using ProcessingOps.CONV (this is a transposed conv) - xt = grad_output - if C.sx > 1 or C.sy > 1: # unstride. NOTE: this is really memory intensive for big strides. (but only when we contiguous it) - xt = xt.movement_op(MovementOps.RESHAPE, (grad_output.shape[0], grad_output.shape[1], grad_output.shape[2], 1, grad_output.shape[3], 1)) - xt = xt.movement_op(MovementOps.PAD, ((0,0), (0,0), (0,0), (0,C.sy-1), (0,0), (0,C.sx-1))) - xt = xt.movement_op(MovementOps.RESHAPE, (xt.shape[0], xt.shape[1], xt.shape[2]*C.sy, xt.shape[4]*C.sx)) - wt = w.movement_op(MovementOps.RESHAPE, (C.groups, C.rcout, C.cin, C.H, C.W)).movement_op(MovementOps.PERMUTE, (0, 2, 1, 3, 4)).movement_op(MovementOps.FLIP, (3, 4)) - wt = wt.movement_op(MovementOps.RESHAPE, (C.groups*C.cin, C.rcout, C.H, C.W)) - py, px = (C.H-1)*C.dy - C.py, (C.W-1)*C.dx - C.px - Cdx = get_conv_args(xt.shape, wt.shape, out_shape=x.shape, dilation=(C.dy, C.dx), padding=(py, px), groups=C.groups) - dx = xt.processing_op(ProcessingOps.CONV, wt, Cdx) - - if self.needs_input_grad[1]: # compute derivative of weights using ProcessingOps.CONV - xdw = x.movement_op(MovementOps.RESHAPE, (C.bs, C.groups, C.cin, C.iy, C.ix)).movement_op(MovementOps.PERMUTE, (2, 1, 0, 3, 4)) - xdw = xdw.movement_op(MovementOps.RESHAPE, (C.cin, C.groups*C.bs, C.iy, C.ix)) - grad_output_dw = grad_output.movement_op(MovementOps.PERMUTE, (1,0,2,3)) - Cdw = get_conv_args(xdw.shape, grad_output_dw.shape, out_shape=(w.shape[1], w.shape[0], w.shape[2], w.shape[3]), padding=(C.py, C.px), stride=(C.dy, C.dx), dilation=(C.sy, C.sx), groups=C.groups) - dw = xdw.processing_op(ProcessingOps.CONV, grad_output_dw, Cdw).movement_op(MovementOps.PERMUTE, (1,0,2,3)) - - return dx, dw diff --git a/tinygrad/ops.py b/tinygrad/ops.py index b1e055885a..98ac10a0ad 100644 --- a/tinygrad/ops.py +++ b/tinygrad/ops.py @@ -11,13 +11,12 @@ from tinygrad.shape import ShapeTracker 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(); STRIDED = auto(); PAD = auto(); SHRINK = 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 ProcessingOps(Enum): CONV = auto() # noqa: E702 class LoadOps(Enum): FROMCPU = auto(); CONTIGUOUS = auto() # noqa: E702 -Op = Union[UnaryOps, BinaryOps, ReduceOps, MovementOps, ProcessingOps, LoadOps, FusedOps] -OpType = Union[Type[UnaryOps], Type[BinaryOps], Type[ReduceOps], Type[MovementOps], Type[ProcessingOps], Type[LoadOps], Type[FusedOps]] +Op = Union[UnaryOps, BinaryOps, ReduceOps, MovementOps, LoadOps, FusedOps] +OpType = Union[Type[UnaryOps], Type[BinaryOps], Type[ReduceOps], Type[MovementOps], Type[LoadOps], Type[FusedOps]] class LazyOp(NamedTuple): op: Op @@ -51,9 +50,7 @@ shape_fxn_for_op : Dict[Op, Callable] = { **{op:lambda self: GenericShape(self.shape, self.flops + prod(self.shape)) for op in UnaryOps}, **{op:lambda self,y: GenericShape(self.shape, self.flops + y.flops + prod(self.shape)) for op in BinaryOps}, **{op:lambda self,new_shape: GenericShape(new_shape, self.flops + prod(self.shape)) for op in ReduceOps}, - **{op:functools.partial(lambda mop,self,arg: GenericShape(ShapeTracker(self.shape).movement_op(mop, arg).shape, self.flops), op) for op in MovementOps}, - # https://docs.nvidia.com/deeplearning/performance/dl-performance-convolutional/index.html - ProcessingOps.CONV:lambda self,w,C: GenericShape(C.out_shape, 2 * (C.bs * C.cout * C.oy * C.ox) * (C.cin * C.H * C.W))} + **{op:functools.partial(lambda mop,self,arg: GenericShape(ShapeTracker(self.shape).movement_op(mop, arg).shape, self.flops), op) for op in MovementOps}} # used in CPUBuffer and TorchBuffer class GenericExecAST(DeviceBuffer): # pylint: disable=abstract-method diff --git a/tinygrad/shape/__init__.py b/tinygrad/shape/__init__.py index d15c9a7d50..adbb2d09d6 100644 --- a/tinygrad/shape/__init__.py +++ b/tinygrad/shape/__init__.py @@ -144,18 +144,6 @@ class ShapeTracker: def needs_valid(self) -> bool: return any(isinstance(v, ZeroView) for v in self.views) - # TODO: do we really need this for conv? - # if we replace, confirm the ops taken fold into one view - def strided(self, arg : Tuple[Tuple[int, int], ...]) -> ShapeTracker: - assert isinstance(arg, tuple) - view = View(tuple(x[0] for x in arg), tuple(x[1] for x in arg)) - # TODO: this does not always require a new view if non contiguous - if self.views[-1].contiguous: - self.views[-1] = view - else: - self.views.append(view) - return self - def reshape(self, new_shape : Tuple[int, ...]) -> ShapeTracker: assert isinstance(new_shape, tuple) if self.shape == new_shape: return self diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 36c26279ca..bc7865151f 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -5,9 +5,6 @@ import numpy as np from typing import List, Tuple, Callable, Optional, ClassVar, Type, Union from tinygrad.helpers import prod, argfix, make_pair, getenv, DEBUG, flatten from tinygrad.lazy import Device, LazyBuffer - -HLOP = getenv("HLOP", 1) - from tinygrad.image import image_conv2d_decorator # An instantiation of the Function is the Context @@ -321,11 +318,6 @@ class Tensor: assert cin*groups == cin_, f"Input Tensor shape {self.shape} does not match the shape of the weights {weight.shape}. ({cin*groups} vs. {cin_})" padding_ = [padding]*4 if isinstance(padding, int) else (padding if len(padding) == 4 else [padding[1], padding[1], padding[0], padding[0]]) - # old implementation - if not HLOP: - ret = mlops.Conv2D.apply(self, weight, groups=groups, stride=stride, dilation=dilation, padding=padding) - return ret if bias is None else ret.add(bias.reshape(1, -1, 1, 1)) - # conv2d is a pooling op (with padding) x = self.pad2d(padding_)._pool((H,W),stride, dilation)