movement ops, reshape is a copy now
MODIFY
README.md
README.md
@@ -137,9 +137,7 @@ Buffer # class of memory on this devic
unary_op (RELU, EXP, LOG, NEG, SIGN) # A -> A
reduce_op (SUM, MAX) # A -> B (smaller size, B has 1 in shape)
binary_op (ADD, SUB, MUL, DIV, POW, A, CMPEQ) # A + B -> C (broadcasting supported)
-reshape # A -> B (different size, same dim count)
-perm_axis # A -> A
-inner_slice # A -> B (different size)
+movement_op (RESHAPE, PERMUTE, SLICE) # A -> B (different size)
conv, convdw, convdx # A + B -> C
```

MODIFY
tinygrad/llops/ops_cpu.py
tinygrad/llops/ops_cpu.py
import numpy as np
-from tinygrad.helpers import get_conv_args, UnaryOps, BinaryOps, ReduceOps
+from tinygrad.helpers import get_conv_args, UnaryOps, BinaryOps, ReduceOps, MovementOps
class CPUBuffer(np.ndarray):
def relu(x): return np.maximum(x, 0)
@@ -48,20 +48,15 @@ def reduce_op(op, inp, ret):
else: raise Exception(f"{op} isn't supported")
return ret
-def reshape(x, ret):
- assert np.prod(x.shape) == np.prod(ret.shape)
- ret[:] = x.reshape(ret.shape)
- return ret
-
-def perm_axis(x, order, ret):
- ret[:] = x.permute(order)
- return ret
-
-def inner_slice(x, arg, ret):
+def movement_op(op, x, ret, arg=None):
+ if op == MovementOps.RESHAPE: ret[:] = x.reshape(ret.shape)
+ elif op == MovementOps.PERMUTE: ret[:] = x.permute(arg)
+ elif op == MovementOps.SLICE:
padding = [(max(0, -p[0]), max(0, p[1]-x.shape[i])) for i,p in enumerate(arg)]
x = x.custompad(padding)
slicee = [(p[0] + padding[i][0], p[1] + padding[i][0]) for i,p in enumerate(arg)]
ret[:] = x[tuple([slice(x[0], x[1], None) for x in slicee])]
+ else: raise Exception(f"{op} isn't supported")
return ret
def get_tx(x, C):

MODIFY
tinygrad/llops/ops_gpu.py
tinygrad/llops/ops_gpu.py
import functools
import numpy as np
import pyopencl as cl
-from tinygrad.helpers import binary_broadcast, get_conv_args, UnaryOps, BinaryOps, ReduceOps
+from tinygrad.helpers import binary_broadcast, get_conv_args, UnaryOps, BinaryOps, ReduceOps, MovementOps
cl_ctx, cl_queue = None, None
def require_init_gpu():
@@ -143,9 +143,8 @@ def reduce_op(op, inp, ret):
buffer_np(np.array(ret.shape, dtype=np.int32)))
return ret
-def reshape(x, shape):
- assert np.prod(x.shape) == np.prod(shape)
- return GPUBuffer(shape, hostbuf=x)
+def reshape(x, ret):
+ cl.enqueue_copy(cl_queue, ret.cl, x.cl)
def perm_axis(inp, order, ret):
perm = clbuild("perm", """
@@ -192,6 +191,12 @@ def inner_slice(x, arg, ret):
buffer_np(np.array(shift, dtype=np.int32)))
return ret
+def movement_op(op, x, ret, arg=None):
+ if op == MovementOps.RESHAPE: reshape(x, ret)
+ elif op == MovementOps.PERMUTE: perm_axis(x, arg, ret)
+ elif op == MovementOps.SLICE: inner_slice(x, arg, ret)
+ return ret
+
def conv(x,w,ret,stride,groups):
C = get_conv_args(x.shape, w.shape, stride, groups)
# input = (bs, groups, cin, iy, ix)

MODIFY
tinygrad/llops/ops_torch.py
tinygrad/llops/ops_torch.py
@@ -20,7 +20,7 @@ class TorchBuffer(torch.Tensor):
# ************* unary+binary+reduce+movement ops *************
-from tinygrad.llops.ops_cpu import unary_op, binary_op, reduce_op, reshape, perm_axis, inner_slice
+from tinygrad.llops.ops_cpu import unary_op, binary_op, reduce_op, movement_op
# ************* processing ops *************

MODIFY
tinygrad/mlops.py
tinygrad/mlops.py
import numpy as np # TODO: remove this, it's used for np.prod and np.argsort
-from tinygrad.helpers import binary_broadcast, get_conv_args, UnaryOps, BinaryOps, ReduceOps
+from tinygrad.helpers import binary_broadcast, get_conv_args, UnaryOps, BinaryOps, ReduceOps, MovementOps
from tinygrad.tensor import Function
# ************* unary ops *************
@@ -126,35 +126,35 @@ class Reshape(Function):
def forward(ctx, x, shape):
ctx.save_for_backward(x.shape)
shape = tuple(-np.prod(x.shape) // np.prod(shape) if s == -1 else s for s in shape)
- return ctx.op.reshape(x, ctx.buffer(shape))
+ return ctx.op.movement_op(MovementOps.RESHAPE, x, ctx.buffer(shape))
def backward(ctx, grad_output):
in_shape, = ctx.saved_tensors
- return ctx.op.reshape(grad_output, ctx.buffer(in_shape))
+ return ctx.op.movement_op(MovementOps.RESHAPE, grad_output, ctx.buffer(in_shape))
class Permute(Function):
def forward(ctx, x, order=(1,0)):
ctx.save_for_backward(order)
ret = ctx.buffer([x.shape[i] for i in order])
- return ctx.op.perm_axis(x, order, ret)
+ return ctx.op.movement_op(MovementOps.PERMUTE, x, ret, order)
def backward(ctx, grad_output):
order, = ctx.saved_tensors
norder = np.argsort(order).tolist()
ret = ctx.buffer([grad_output.shape[i] for i in norder])
- return ctx.op.perm_axis(grad_output, norder, ret)
+ return ctx.op.movement_op(MovementOps.PERMUTE, grad_output, ret, norder)
class Slice(Function):
def forward(ctx, x, arg=None):
ctx.save_for_backward(x.shape, arg)
ret = ctx.buffer([y[1]-y[0] for y in arg])
- return ctx.op.inner_slice(x, arg, ret)
+ return ctx.op.movement_op(MovementOps.SLICE, x, ret, arg)
def backward(ctx, grad_output):
shape, arg = ctx.saved_tensors
narg = [(0-p[0], grad_output.shape[i]+(shape[i]-p[1])) for i,p in enumerate(arg)]
ret = ctx.buffer([y[1]-y[0] for y in narg])
- return ctx.op.inner_slice(grad_output, narg, ret)
+ return ctx.op.movement_op(MovementOps.SLICE, grad_output, ret, narg)
# ************* processing ops *************

movement ops, reshape is a copy now