forked from tinygrad/tinygrad
George Hotz
+2
-3
@@ -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):
|
||||
|
||||
@@ -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
|
||||
|
||||
|
||||
+3
-3
@@ -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:
|
||||
|
||||
+1
-25
@@ -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))
|
||||
|
||||
+2
-5
@@ -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
|
||||
|
||||
+9
-114
@@ -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
|
||||
|
||||
|
||||
@@ -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):
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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))
|
||||
}}
|
||||
|
||||
|
||||
+2
-36
@@ -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
|
||||
|
||||
+4
-7
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user