mlops for cpu
MODIFY
tinygrad/llops/cpu.py
tinygrad/llops/cpu.py
import numpy as np
from tinygrad.helpers import UnaryOps, BinaryOps, ReduceOps
-class CPUBuffer(np.ndarray):
+class Buffer(np.ndarray):
def __new__(cls, shape, hostbuf=None):
if hostbuf is not None:
return super().__new__(cls, shape, dtype=np.float32, buffer=hostbuf.data)

MODIFY
tinygrad/llops/opencl.py
tinygrad/llops/opencl.py
@@ -18,12 +18,12 @@ i32 = np.int32
def roundup(x, n=4): return (x+(n-1))//n * n
def sync(): cl_queue.finish()
-class GPUBuffer:
+class Buffer:
def __init__(self, shape, hostbuf=None):
require_init_gpu()
self.shape, self.dtype = tuple(shape), np.float32
- self.cl = hostbuf.cl if isinstance(hostbuf, GPUBuffer) else cl.Buffer(cl_ctx, cl.mem_flags.READ_WRITE, 4*roundup(np.prod(shape))) # padding
- if hostbuf is not None and not isinstance(hostbuf, GPUBuffer):
+ self.cl = hostbuf.cl if isinstance(hostbuf, Buffer) else cl.Buffer(cl_ctx, cl.mem_flags.READ_WRITE, 4*roundup(np.prod(shape))) # padding
+ if hostbuf is not None and not isinstance(hostbuf, Buffer):
cl.enqueue_copy(cl_queue, self.cl, hostbuf.astype(np.float32).ravel())
def __repr__(self):
@@ -31,7 +31,7 @@ class GPUBuffer:
@staticmethod
def fromCPU(x):
- return GPUBuffer(x.shape, x.view(np.ndarray))
+ return Buffer(x.shape, x.view(np.ndarray))
def toCPU(self):
data = np.empty(self.shape, dtype=np.float32)
@@ -316,4 +316,3 @@ def convdx(w,grad_output,dx,conv_args):
""")
convdx_prg([bs, groups, cin], None, w.cl, grad_output.cl, dx.cl, *[i32(x) for x in conv_args])
return dx
-

ADD
tinygrad/mlops.py
+import numpy as np
+from tinygrad.helpers import binary_broadcast, UnaryOps, BinaryOps, ReduceOps
+from tinygrad.tensor import Function
+
+ll = None
+def select_llops(ops):
+ global ll
+ import importlib
+ ll = importlib.import_module(f".{ops}", f"tinygrad.llops")
+ return ll.Buffer
+
+# ************* unary ops *************
+
+class UnaryOp(Function):
+ def forward(ctx, input):
+ ctx.save_for_backward(input)
+ return ll.unary_op(ctx.fop, input, ll.Buffer(input.shape))
+
+ def backward(ctx, grad_output):
+ input, = ctx.saved_tensors
+ return ll.binary_op(ctx.bop, input, grad_output, ll.Buffer(input.shape))
+
+class ReLU(UnaryOp):
+ fop = UnaryOps.RELU
+
+ def backward(ctx, grad_output):
+ input, = ctx.saved_tensors
+ ret = ll.Buffer(input.shape)
+ ll.unary_op(UnaryOps.SIGN, input, ret)
+ ll.unary_op(UnaryOps.RELU, ret, ret)
+ return ll.binary_op(BinaryOps.MUL, ret, grad_output, ret)
+
+class Log(UnaryOp):
+ fop = UnaryOps.LOG
+ bop = BinaryOps.DIV
+
+class Exp(UnaryOp):
+ def forward(ctx, input):
+ ret = ll.unary_op(UnaryOps.EXP, input, ll.Buffer(input.shape))
+ ctx.save_for_backward(ret) # we save the output here, not the input
+ return ret
+
+ bop = BinaryOps.MUL
+
+# ************* reduce ops *************
+
+def reduce_shape(shape, axis):
+ return [1 if i in axis else shape[i] for i in range(len(shape))]
+
+class Sum(Function):
+ def forward(ctx, input, axis=None):
+ ctx.save_for_backward(input.shape)
+ return ll.reduce_op(ReduceOps.SUM, input, ll.Buffer(reduce_shape(input.shape, axis)))
+
+ def backward(ctx, grad_output):
+ shape_input, = ctx.saved_tensors
+ # NOTE: the b Buffer isn't used, since this is just for broadcast
+ ret = ll.Buffer(shape_input)
+ return ll.binary_op(BinaryOps.A, grad_output, ret, ret)
+
+class Max(Function):
+ def forward(ctx, input, axis=None):
+ ret = ll.reduce_op(ReduceOps.MAX, input, ll.Buffer(reduce_shape(input.shape, axis)))
+ ctx.save_for_backward(input, ret)
+ return ret
+
+ def backward(ctx, grad_output):
+ input, ret = ctx.saved_tensors
+ ret2 = ll.binary_op(BinaryOps.CMPEQ, input, ret, ll.Buffer(input.shape))
+ div = ll.reduce_op(ReduceOps.SUM, ret2, ll.Buffer(grad_output.shape))
+ ll.binary_op(BinaryOps.DIV, div, ret2, ret2)
+ return ll.binary_op(BinaryOps.MUL, ret2, grad_output, ret2)
+
+# ************* binary ops *************
+
+def unbroadcast(out, in_sh):
+ return ll.reduce_op(ReduceOps.SUM, out, ll.Buffer(in_sh))
+
+class Add(Function):
+ def forward(ctx, x, y):
+ ctx.save_for_backward(x.shape, y.shape)
+ return ll.binary_op(BinaryOps.ADD, x, y, ll.Buffer(binary_broadcast(x.shape, y.shape)))
+
+ def backward(ctx, grad_output):
+ shape_x, shape_y = ctx.saved_tensors
+ return unbroadcast(grad_output, shape_x) if ctx.needs_input_grad[0] else None, \
+ unbroadcast(grad_output, shape_y) if ctx.needs_input_grad[1] else None
+
+class Sub(Function):
+ def forward(ctx, x, y):
+ ctx.save_for_backward(x.shape, y.shape)
+ return ll.binary_op(BinaryOps.SUB, x, y, ll.Buffer(binary_broadcast(x.shape, y.shape)))
+
+ def backward(ctx, grad_output):
+ shape_x, shape_y = ctx.saved_tensors
+ neg_grad_output = ll.unary_op(UnaryOps.NEG, grad_output, ll.Buffer(grad_output.shape))
+ return unbroadcast(grad_output, shape_x) if ctx.needs_input_grad[0] else None, \
+ unbroadcast(neg_grad_output, shape_y) if ctx.needs_input_grad[1] else None
+
+class Mul(Function):
+ def forward(ctx, x, y):
+ ctx.save_for_backward(x, y)
+ return ll.binary_op(BinaryOps.MUL, x, y, ll.Buffer(binary_broadcast(x.shape, y.shape)))
+
+ def backward(ctx, grad_output):
+ x,y = ctx.saved_tensors
+ tmp = ll.Buffer(grad_output.shape)
+ grad_x = unbroadcast(ll.binary_op(BinaryOps.MUL, y, grad_output, tmp), x.shape) if ctx.needs_input_grad[0] else None
+ grad_y = unbroadcast(ll.binary_op(BinaryOps.MUL, x, grad_output, tmp), y.shape) if ctx.needs_input_grad[1] else None
+ return grad_x, grad_y
+
+class Pow(Function):
+ def forward(ctx, x, y):
+ ret = ll.Buffer(binary_broadcast(x.shape, y.shape))
+ ctx.save_for_backward(x, y, ret)
+ return ll.binary_op(BinaryOps.POW, x, y, ret)
+
+ def backward(ctx, grad_output):
+ x,y,powxy = ctx.saved_tensors
+ tmp = ll.Buffer(grad_output.shape)
+ ll.binary_op(BinaryOps.DIV, x, powxy, tmp) # pow(x,y)/x
+ ll.binary_op(BinaryOps.MUL, y, tmp, tmp) # y * pow(x,y)/x
+ grad_x = unbroadcast(ll.binary_op(BinaryOps.MUL, grad_output, tmp, tmp), x.shape) if ctx.needs_input_grad[0] else None
+ log_x = ll.unary_op(UnaryOps.LOG, x, ll.Buffer(x.shape))
+ ll.binary_op(BinaryOps.MUL, log_x, powxy, tmp) # log(x) * pow(x,y)
+ grad_y = unbroadcast(ll.binary_op(BinaryOps.MUL, grad_output, tmp, tmp), y.shape) if ctx.needs_input_grad[1] else None
+ return grad_x, grad_y
+
+# ************* movement ops *************
+
+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)
+ r = ll.Buffer(shape, hostbuf=x) # NOTE: this is not a copy
+ assert np.prod(x.shape) == np.prod(r.shape)
+ return r
+
+ def backward(ctx, grad_output):
+ in_shape, = ctx.saved_tensors
+ return ll.Buffer(in_shape, hostbuf=grad_output)
+
+class Transpose(Function):
+ def forward(ctx, x, order=(1,0)):
+ ctx.save_for_backward(order)
+ ret = ll.Buffer([x.shape[i] for i in order])
+ return ll.perm_axis(x, order, ret)
+
+ def backward(ctx, grad_output):
+ norder = np.argsort(ctx.order)
+ ret = ll.Buffer([grad_output.shape[i] for i in norder])
+ return ll.perm_axis(grad_output, norder, ret)
+
+class Slice(Function):
+ def forward(ctx, x, arg=None):
+ ctx.save_for_backward(x.shape)
+ ret = ll.Buffer([y[1]-y[0] for y in arg])
+ return ll.inner_slice(x, arg, ret)
+
+ def backward(ctx, grad_output):
+ shape, = ctx.saved_tensors
+ narg = [(0-p[0], grad_output.shape[i]+(shape[i]-p[1])) for i,p in enumerate(ctx.arg)]
+ ret = ll.Buffer([y[1]-y[0] for y in narg])
+ return ll.inner_slice(grad_output, narg, ret)
+
+# ************* processing ops *************
+
+class Matmul(Function):
+ def forward(ctx, input, weight):
+ assert input.shape[-1] == weight.shape[-2]
+ ret = ll.Buffer(list(input.shape[0:-1])+[weight.shape[-1]])
+ ctx.save_for_backward(input, weight)
+ return ll.matmul(input, weight, ret)
+
+ def backward(ctx, grad_output):
+ input, weight = ctx.saved_tensors
+ grad_input = ll.matmul(grad_output, weight, ll.Buffer(input.shape), transpose_b=True) if ctx.needs_input_grad[0] else None
+ grad_weight = ll.matmul(input, grad_output, ll.Buffer(weight.shape), transpose_a=True) if ctx.needs_input_grad[1] else None
+ return grad_input, grad_weight
+
+class Conv2D(Function):
+ def forward(ctx, x, w, stride=1, groups=1):
+ if isinstance(ctx.stride, int): ctx.stride = (ctx.stride, ctx.stride)
+ cout,cin,H,W = w.shape
+ ys,xs = ctx.stride
+ bs,cin_,iy,ix = x.shape
+ oy,ox = (iy-(H-ys))//ys, (ix-(W-xs))//xs
+ if cin*ctx.groups != cin_: raise Exception(f"Input Tensor shape {x.shape} does not match the shape of the weights {w.shape}. ({cin*ctx.groups} vs. {cin_})")
+ assert cout % ctx.groups == 0
+ rcout = cout//ctx.groups
+
+ ctx.save_for_backward(x,w)
+
+ # output buffer
+ conv_args = H, W, ctx.groups, rcout, cin, oy, ox, iy, ix, ys, xs, bs
+ return ll.conv(x, w, ll.Buffer((bs, cout, oy, ox)), conv_args)
+
+ def backward(ctx, grad_output):
+ bs,_,oy,ox = grad_output.shape
+ x, w = ctx.saved_tensors
+ cout,cin,H,W = w.shape
+ ys,xs = ctx.stride
+ bs,cin_,iy,ix = x.shape
+ oy,ox = (iy-(H-ys))//ys, (ix-(W-xs))//xs
+ assert cin*ctx.groups == cin_
+ assert cout % ctx.groups == 0
+ rcout = cout//ctx.groups
+
+ conv_args = H, W, ctx.groups, rcout, cin, oy, ox, iy, ix, ys, xs, bs
+ dx = ll.convdx(w, grad_output, ll.Buffer((bs, cin_, iy, ix)), conv_args) if ctx.needs_input_grad[0] else None
+ dw = ll.convdw(x, grad_output, ll.Buffer((cout, cin, H, W)), conv_args) if ctx.needs_input_grad[1] else None
+ return dx, dw
\ No newline at end of file

MODIFY
tinygrad/ops/ops_cpu.py
tinygrad/ops/ops_cpu.py
-import numpy as np
-from ..tensor import Function
-
-class CPUBuffer(np.ndarray):
- def toCPU(x): return x
- def log(x): return np.log(x)
- def exp(x): return np.exp(x)
- def sign(x): return np.sign(x)
- def relu(x): return np.maximum(x, 0)
- def type(x, tt): return x.astype(tt)
- def custompad(x, padding): return np.pad(x, padding)
- def permute(x, order): return x.transpose(order)
- def expand(x, shp): return np.broadcast_to(x, shp)
- def amax(x, *args, **kwargs): return np.amax(x, *args, **kwargs)
-
- @staticmethod
- def fromCPU(x): return x
-
-# ************* unary ops *************
-
-class ReLU(Function):
- def forward(ctx, input):
- ctx.save_for_backward(input)
- return input.relu()
-
- def backward(ctx, grad_output):
- input, = ctx.saved_tensors
- return grad_output * (input >= 0)
-
-class Log(Function):
- def forward(ctx, input):
- ctx.save_for_backward(input)
- return input.log()
-
- def backward(ctx, grad_output):
- input, = ctx.saved_tensors
- return grad_output / input
-
-class Exp(Function):
- def forward(ctx, input):
- ret = input.clip(-88, 88).exp()
- ctx.save_for_backward(ret)
- return ret
-
- def backward(ctx, grad_output):
- ret, = ctx.saved_tensors
- return grad_output * ret
-
-# ************* reduce ops (with keepdims=True) *************
-
-class Sum(Function):
- def forward(ctx, input, axis):
- ctx.save_for_backward(input.shape)
- return input.sum(axis, keepdims=True)
-
- def backward(ctx, grad_output):
- shape_input, = ctx.saved_tensors
- return grad_output.expand(shape_input)
-
-class Max(Function):
- def forward(ctx, inp, axis):
- ret = inp.amax(axis=axis, keepdims=True)
- ctx.save_for_backward(inp, axis, ret)
- return ret
-
- def backward(ctx, grad_output):
- input, axis, ret = ctx.saved_tensors
- ret2 = (input==ret)
- div = ret2.sum(axis=tuple(axis), keepdims=True)
- return ret2*grad_output/div.type(input.dtype)
-
-# ************* binary ops (with broadcasting) *************
-
-def unbroadcast(out, in_sh):
- # adjoint operation to broadcast is sum. Need to sum all axis with 1 = in_sh[i] < out.shape[i]
- if in_sh == (1,):
- return out.sum().reshape((1,))
- else:
- sum_axis = tuple([i for i in range(len(in_sh)) if in_sh[i]==1 and out.shape[i]>1])
- return out.sum(axis=sum_axis).reshape(in_sh) if len(sum_axis) > 0 else out
-
-class Add(Function):
- def forward(ctx, x, y):
- ctx.save_for_backward(x.shape, y.shape)
- return x+y
-
- def backward(ctx, grad_output):
- shape_x, shape_y = ctx.saved_tensors
- return unbroadcast(grad_output, shape_x), unbroadcast(grad_output, shape_y)
-
-class Sub(Function):
- def forward(ctx, x, y):
- ctx.save_for_backward(x.shape, y.shape)
- return x-y
-
- def backward(ctx, grad_output):
- shape_x, shape_y = ctx.saved_tensors
- return unbroadcast(grad_output, shape_x), unbroadcast(-grad_output, shape_y)
-
-class Mul(Function):
- def forward(ctx, x, y):
- ctx.save_for_backward(x, y)
- return x*y
-
- def backward(ctx, grad_output):
- x,y = ctx.saved_tensors
- return unbroadcast(y*grad_output, x.shape), unbroadcast(x*grad_output, y.shape)
-
-class Pow(Function):
- def forward(ctx, x, y):
- ctx.save_for_backward(x, y)
- return x ** y
-
- def backward(ctx, grad_output):
- x,y = ctx.saved_tensors
- return unbroadcast(y * (x**(y-1.0)) * grad_output, x.shape), \
- unbroadcast((x**y) * x.log() * grad_output, y.shape)
-
-# ************* movement ops *************
-
-class Reshape(Function):
- def forward(ctx, x, shape):
- ctx.save_for_backward(x.shape)
- return x.reshape(shape)
-
- def backward(ctx, grad_output):
- in_shape, = ctx.saved_tensors
- return grad_output.reshape(in_shape)
-
-class Transpose(Function):
- def forward(ctx, x, order):
- ctx.save_for_backward(order)
- return x.permute(order)
-
- def backward(ctx, x):
- return x.permute(tuple(np.argsort(ctx.order)))
-
-def inner_slice(x, arg):
- 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)]
- return x[tuple([slice(x[0], x[1], None) for x in slicee])]
-
-class Slice(Function):
- def forward(ctx, x, arg=None):
- ctx.save_for_backward(x.shape)
- return inner_slice(x, arg)
-
- def backward(ctx, grad_output):
- shape, = ctx.saved_tensors
- narg = [(0-p[0], grad_output.shape[i]+(shape[i]-p[1])) for i,p in enumerate(ctx.arg)]
- return inner_slice(grad_output, narg)
-
-# ************* processing ops *************
-
-class Matmul(Function):
- def forward(ctx, input, weight):
- ctx.save_for_backward(input, weight)
- return input @ weight
-
- def backward(ctx, grad_output):
- input, weight = ctx.saved_tensors
- grad_input = grad_output @ weight.swapaxes(-2, -1)
- grad_weight = input.swapaxes(-2, -1) @ grad_output
- return grad_input, grad_weight
-
-class Conv2D(Function):
- def forward(ctx, x, w, stride=1, groups=1):
- if isinstance(ctx.stride, int): ctx.stride = (ctx.stride, ctx.stride)
- cout,cin,H,W = w.shape
- ys,xs = ctx.stride
- bs,cin_ = x.shape[0], x.shape[1]
- oy,ox = (x.shape[2]-(H-ys))//ys, (x.shape[3]-(W-xs))//xs
- assert cin*ctx.groups == cin_
- assert cout % ctx.groups == 0
- rcout = cout//ctx.groups
-
- gx = x.reshape(bs,ctx.groups,cin,x.shape[2],x.shape[3])
- tx = np.lib.stride_tricks.as_strided(gx,
- shape=(bs, ctx.groups, cin, oy, ox, H, W),
- strides=(*gx.strides[0:3], gx.strides[3]*ys, gx.strides[4]*xs, *gx.strides[3:5]),
- writeable=False,
- )
- tw = w.reshape(ctx.groups, rcout, cin, H, W)
- ctx.save_for_backward(tx, tw, x.shape)
-
- ret = np.zeros((bs,ctx.groups,oy,ox,rcout),dtype=x.dtype)
- for g in range(ctx.groups):
- #ijYXyx,kjyx -> iYXk ->ikYX
- ret[:,g] += np.tensordot(tx[:,g], tw[g], ((1,4,5),(1,2,3)))
- return np.moveaxis(ret,4,2).reshape(bs, cout, oy, ox)
-
- def backward(ctx, grad_output):
- bs,_,oy,ox = grad_output.shape
- tx, tw, x_shape = ctx.saved_tensors
- _,rcout,cin,H,W = tw.shape
- ys,xs = ctx.stride
- OY,OX = x_shape[2:4]
-
- ggg = grad_output.reshape(bs,ctx.groups,rcout,oy,ox)
-
- gdw = np.zeros((ctx.groups,rcout,cin,H,W), dtype=tx.dtype)
- for g in range(ctx.groups):
- #'ikYX,ijYXyx -> kjyx'
- gdw[g] += np.tensordot(ggg[:,g], tx[:,g], ((0,2,3),(0,2,3)))
-
- # needs to be optimized
- gdx = np.zeros((bs,ctx.groups,cin,OY,OX), dtype=tx.dtype)
- for k in range(oy*ox):
- Y, X = k//ox, k%ox
- iY,iX = Y*ys, X*xs
- #gdx[:,:,: , iY:iY+H, iX:iX+W] += np.einsum('igk,gkjyx->igjyx', ggg[:,:,:,Y,X], tw)
- for g in range(ctx.groups):
- tg = np.dot(ggg[:,g,:,Y,X].reshape(bs, -1), tw[g].reshape(rcout, -1))
- gdx[:, g, :, iY:iY+H, iX:iX+W] += tg.reshape((bs, cin, H, W))
-
- return gdx.reshape((bs, ctx.groups*cin, OY, OX)), gdw.reshape((ctx.groups*rcout, cin, H, W))
+from ..mlops import *
+Buffer = select_llops("cpu")
\ No newline at end of file

MODIFY
tinygrad/ops/ops_gpu.py
tinygrad/ops/ops_gpu.py
-import os
-import numpy as np
-from tinygrad.helpers import binary_broadcast, UnaryOps, BinaryOps, ReduceOps
-from ..tensor import Function
-
-if int(os.getenv("LLCPU", 0)) == 1:
- from ..llops.cpu import CPUBuffer as Buffer
- from ..llops.cpu import unary_op, binary_op, reduce_op, perm_axis, inner_slice
- from ..llops.cpu import matmul, conv, convdw, convdx
-else:
- from ..llops.opencl import GPUBuffer as Buffer
- from ..llops.opencl import unary_op, binary_op, reduce_op, perm_axis, inner_slice
- from ..llops.opencl import matmul, conv, convdw, convdx
-
-# ************* unary ops *************
-
-class UnaryOp(Function):
- def forward(ctx, input):
- ctx.save_for_backward(input)
- return unary_op(ctx.fop, input, Buffer(input.shape))
-
- def backward(ctx, grad_output):
- input, = ctx.saved_tensors
- return binary_op(ctx.bop, input, grad_output, Buffer(input.shape))
-
-class ReLU(UnaryOp):
- fop = UnaryOps.RELU
-
- def backward(ctx, grad_output):
- input, = ctx.saved_tensors
- ret = Buffer(input.shape)
- unary_op(UnaryOps.SIGN, input, ret)
- unary_op(UnaryOps.RELU, ret, ret)
- return binary_op(BinaryOps.MUL, ret, grad_output, ret)
-
-class Log(UnaryOp):
- fop = UnaryOps.LOG
- bop = BinaryOps.DIV
-
-class Exp(UnaryOp):
- def forward(ctx, input):
- ret = unary_op(UnaryOps.EXP, input, Buffer(input.shape))
- ctx.save_for_backward(ret) # we save the output here, not the input
- return ret
-
- bop = BinaryOps.MUL
-
-# ************* reduce ops *************
-
-def reduce_shape(shape, axis):
- return [1 if i in axis else shape[i] for i in range(len(shape))]
-
-class Sum(Function):
- def forward(ctx, input, axis=None):
- ctx.save_for_backward(input.shape)
- return reduce_op(ReduceOps.SUM, input, Buffer(reduce_shape(input.shape, axis)))
-
- def backward(ctx, grad_output):
- shape_input, = ctx.saved_tensors
- # NOTE: the b Buffer isn't used, since this is just for broadcast
- ret = Buffer(shape_input)
- return binary_op(BinaryOps.A, grad_output, ret, ret)
-
-class Max(Function):
- def forward(ctx, input, axis=None):
- ret = reduce_op(ReduceOps.MAX, input, Buffer(reduce_shape(input.shape, axis)))
- ctx.save_for_backward(input, ret)
- return ret
-
- def backward(ctx, grad_output):
- input, ret = ctx.saved_tensors
- ret2 = binary_op(BinaryOps.CMPEQ, input, ret, Buffer(input.shape))
- div = reduce_op(ReduceOps.SUM, ret2, Buffer(grad_output.shape))
- binary_op(BinaryOps.DIV, div, ret2, ret2)
- return binary_op(BinaryOps.MUL, ret2, grad_output, ret2)
-
-# ************* binary ops *************
-
-def unbroadcast(out, in_sh):
- return reduce_op(ReduceOps.SUM, out, Buffer(in_sh))
-
-class Add(Function):
- def forward(ctx, x, y):
- ctx.save_for_backward(x.shape, y.shape)
- return binary_op(BinaryOps.ADD, x, y, Buffer(binary_broadcast(x.shape, y.shape)))
-
- def backward(ctx, grad_output):
- shape_x, shape_y = ctx.saved_tensors
- return unbroadcast(grad_output, shape_x) if ctx.needs_input_grad[0] else None, \
- unbroadcast(grad_output, shape_y) if ctx.needs_input_grad[1] else None
-
-class Sub(Function):
- def forward(ctx, x, y):
- ctx.save_for_backward(x.shape, y.shape)
- return binary_op(BinaryOps.SUB, x, y, Buffer(binary_broadcast(x.shape, y.shape)))
-
- def backward(ctx, grad_output):
- shape_x, shape_y = ctx.saved_tensors
- neg_grad_output = unary_op(UnaryOps.NEG, grad_output, Buffer(grad_output.shape))
- return unbroadcast(grad_output, shape_x) if ctx.needs_input_grad[0] else None, \
- unbroadcast(neg_grad_output, shape_y) if ctx.needs_input_grad[1] else None
-
-class Mul(Function):
- def forward(ctx, x, y):
- ctx.save_for_backward(x, y)
- return binary_op(BinaryOps.MUL, x, y, Buffer(binary_broadcast(x.shape, y.shape)))
-
- def backward(ctx, grad_output):
- x,y = ctx.saved_tensors
- tmp = Buffer(grad_output.shape)
- grad_x = unbroadcast(binary_op(BinaryOps.MUL, y, grad_output, tmp), x.shape) if ctx.needs_input_grad[0] else None
- grad_y = unbroadcast(binary_op(BinaryOps.MUL, x, grad_output, tmp), y.shape) if ctx.needs_input_grad[1] else None
- return grad_x, grad_y
-
-class Pow(Function):
- def forward(ctx, x, y):
- ret = Buffer(binary_broadcast(x.shape, y.shape))
- ctx.save_for_backward(x, y, ret)
- return binary_op(BinaryOps.POW, x, y, ret)
-
- def backward(ctx, grad_output):
- x,y,powxy = ctx.saved_tensors
- tmp = Buffer(grad_output.shape)
- binary_op(BinaryOps.DIV, x, powxy, tmp) # pow(x,y)/x
- binary_op(BinaryOps.MUL, y, tmp, tmp) # y * pow(x,y)/x
- grad_x = unbroadcast(binary_op(BinaryOps.MUL, grad_output, tmp, tmp), x.shape) if ctx.needs_input_grad[0] else None
- log_x = unary_op(UnaryOps.LOG, x, Buffer(x.shape))
- binary_op(BinaryOps.MUL, log_x, powxy, tmp) # log(x) * pow(x,y)
- grad_y = unbroadcast(binary_op(BinaryOps.MUL, grad_output, tmp, tmp), y.shape) if ctx.needs_input_grad[1] else None
- return grad_x, grad_y
-
-# ************* movement ops *************
-
-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)
- r = Buffer(shape, hostbuf=x) # NOTE: this is not a copy
- assert np.prod(x.shape) == np.prod(r.shape)
- return r
-
- def backward(ctx, grad_output):
- in_shape, = ctx.saved_tensors
- return Buffer(in_shape, hostbuf=grad_output)
-
-class Transpose(Function):
- def forward(ctx, x, order=(1,0)):
- ctx.save_for_backward(order)
- ret = Buffer([x.shape[i] for i in order])
- return perm_axis(x, order, ret)
-
- def backward(ctx, grad_output):
- norder = np.argsort(ctx.order)
- ret = Buffer([grad_output.shape[i] for i in norder])
- return perm_axis(grad_output, norder, ret)
-
-class Slice(Function):
- def forward(ctx, x, arg=None):
- ctx.save_for_backward(x.shape)
- ret = Buffer([y[1]-y[0] for y in arg])
- return inner_slice(x, arg, ret)
-
- def backward(ctx, grad_output):
- shape, = ctx.saved_tensors
- narg = [(0-p[0], grad_output.shape[i]+(shape[i]-p[1])) for i,p in enumerate(ctx.arg)]
- ret = Buffer([y[1]-y[0] for y in narg])
- return inner_slice(grad_output, narg, ret)
-
-# ************* processing ops *************
-
-class Matmul(Function):
- def forward(ctx, input, weight):
- assert input.shape[-1] == weight.shape[-2]
- ret = Buffer(list(input.shape[0:-1])+[weight.shape[-1]])
- ctx.save_for_backward(input, weight)
- return matmul(input, weight, ret)
-
- def backward(ctx, grad_output):
- input, weight = ctx.saved_tensors
- grad_input = matmul(grad_output, weight, Buffer(input.shape), transpose_b=True) if ctx.needs_input_grad[0] else None
- grad_weight = matmul(input, grad_output, Buffer(weight.shape), transpose_a=True) if ctx.needs_input_grad[1] else None
- return grad_input, grad_weight
-
-class Conv2D(Function):
- def forward(ctx, x, w, stride=1, groups=1):
- if isinstance(ctx.stride, int): ctx.stride = (ctx.stride, ctx.stride)
- cout,cin,H,W = w.shape
- ys,xs = ctx.stride
- bs,cin_,iy,ix = x.shape
- oy,ox = (iy-(H-ys))//ys, (ix-(W-xs))//xs
- if cin*ctx.groups != cin_: raise Exception(f"Input Tensor shape {x.shape} does not match the shape of the weights {w.shape}. ({cin*ctx.groups} vs. {cin_})")
- assert cout % ctx.groups == 0
- rcout = cout//ctx.groups
-
- ctx.save_for_backward(x,w)
-
- # output buffer
- conv_args = H, W, ctx.groups, rcout, cin, oy, ox, iy, ix, ys, xs, bs
- return conv(x, w, Buffer((bs, cout, oy, ox)), conv_args)
-
- def backward(ctx, grad_output):
- bs,_,oy,ox = grad_output.shape
- x, w = ctx.saved_tensors
- cout,cin,H,W = w.shape
- ys,xs = ctx.stride
- bs,cin_,iy,ix = x.shape
- oy,ox = (iy-(H-ys))//ys, (ix-(W-xs))//xs
- assert cin*ctx.groups == cin_
- assert cout % ctx.groups == 0
- rcout = cout//ctx.groups
-
- conv_args = H, W, ctx.groups, rcout, cin, oy, ox, iy, ix, ys, xs, bs
- dx = convdx(w, grad_output, Buffer((bs, cin_, iy, ix)), conv_args) if ctx.needs_input_grad[0] else None
- dw = convdw(x, grad_output, Buffer((cout, cin, H, W)), conv_args) if ctx.needs_input_grad[1] else None
- return dx, dw
+from ..mlops import *
+#Buffer = select_llops("opencl")

mlops for cpu