Ops layer of indirection
MODIFY
tinygrad/helpers.py
tinygrad/helpers.py
@@ -36,6 +36,7 @@ def get_conv_args(x_shape, w_shape, stride, groups):
rcout = cout//groups
return conv_args(H, W, groups, rcout, cin, oy, ox, iy, ix, ys, xs, bs)
+# TODO: move to ops.py
from enum import Enum
UnaryOps = Enum("UnaryOps", ["RELU", "EXP", "LOG", "NEG", "SIGN"])
BinaryOps = Enum("BinaryOps", ["ADD", "SUB", "MUL", "DIV", "POW", "A", "CMPEQ"])

MODIFY
tinygrad/mlops.py
tinygrad/mlops.py
@@ -7,11 +7,11 @@ from tinygrad.tensor import Function
class _UnaryOp(Function):
def forward(ctx, input):
ctx.save_for_backward(input)
- return ctx.op.unary_op(ctx.fop, input, ctx.buffer(input.shape))
+ return ctx.unary_op(ctx.fop, input, ctx.buffer(input.shape))
def backward(ctx, grad_output):
input, = ctx.saved_tensors
- return ctx.op.binary_op(ctx.bop, input, grad_output, ctx.buffer(input.shape))
+ return ctx.binary_op(ctx.bop, input, grad_output, ctx.buffer(input.shape))
class ReLU(_UnaryOp):
fop = UnaryOps.RELU
@@ -19,9 +19,9 @@ class ReLU(_UnaryOp):
def backward(ctx, grad_output):
input, = ctx.saved_tensors
ret = ctx.buffer(input.shape)
- ctx.op.unary_op(UnaryOps.SIGN, input, ret)
- ctx.op.unary_op(UnaryOps.RELU, ret, ret)
- return ctx.op.binary_op(BinaryOps.MUL, ret, grad_output, ret)
+ ctx.unary_op(UnaryOps.SIGN, input, ret)
+ ctx.unary_op(UnaryOps.RELU, ret, ret)
+ return ctx.binary_op(BinaryOps.MUL, ret, grad_output, ret)
class Log(_UnaryOp):
fop = UnaryOps.LOG
@@ -29,7 +29,7 @@ class Log(_UnaryOp):
class Exp(_UnaryOp):
def forward(ctx, input):
- ret = ctx.op.unary_op(UnaryOps.EXP, input, ctx.buffer(input.shape))
+ ret = ctx.unary_op(UnaryOps.EXP, input, ctx.buffer(input.shape))
ctx.save_for_backward(ret) # we save the output here, not the input
return ret
@@ -43,37 +43,37 @@ def reduce_shape(shape, axis):
class Sum(Function):
def forward(ctx, input, axis=None):
ctx.save_for_backward(input.shape)
- return ctx.op.reduce_op(ReduceOps.SUM, input, ctx.buffer(reduce_shape(input.shape, axis)))
+ return ctx.reduce_op(ReduceOps.SUM, input, ctx.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 = ctx.buffer(shape_input)
- return ctx.op.binary_op(BinaryOps.A, grad_output, ret, ret)
+ return ctx.binary_op(BinaryOps.A, grad_output, ret, ret)
class Max(Function):
def forward(ctx, input, axis=None):
- ret = ctx.op.reduce_op(ReduceOps.MAX, input, ctx.buffer(reduce_shape(input.shape, axis)))
+ ret = ctx.reduce_op(ReduceOps.MAX, input, ctx.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 = ctx.op.binary_op(BinaryOps.CMPEQ, input, ret, ctx.buffer(input.shape))
- div = ctx.op.reduce_op(ReduceOps.SUM, ret2, ctx.buffer(grad_output.shape))
- ctx.op.binary_op(BinaryOps.DIV, div, ret2, ret2)
- return ctx.op.binary_op(BinaryOps.MUL, ret2, grad_output, ret2)
+ ret2 = ctx.binary_op(BinaryOps.CMPEQ, input, ret, ctx.buffer(input.shape))
+ div = ctx.reduce_op(ReduceOps.SUM, ret2, ctx.buffer(grad_output.shape))
+ ctx.binary_op(BinaryOps.DIV, div, ret2, ret2)
+ return ctx.binary_op(BinaryOps.MUL, ret2, grad_output, ret2)
# ************* binary ops *************
def unbroadcast(ctx, out, in_sh):
- return ctx.op.reduce_op(ReduceOps.SUM, out, ctx.buffer(in_sh))
+ return ctx.reduce_op(ReduceOps.SUM, out, ctx.buffer(in_sh))
class Add(Function):
def forward(ctx, x, y):
ctx.save_for_backward(x.shape, y.shape)
buf = ctx.buffer(binary_broadcast(x.shape, y.shape))
- return ctx.op.binary_op(BinaryOps.ADD, x, y, buf) #ctx.buffer(binary_broadcast(x.shape, y.shape)))
+ return ctx.binary_op(BinaryOps.ADD, x, y, buf) #ctx.buffer(binary_broadcast(x.shape, y.shape)))
def backward(ctx, grad_output):
shape_x, shape_y = ctx.saved_tensors
@@ -83,41 +83,41 @@ class Add(Function):
class Sub(Function):
def forward(ctx, x, y):
ctx.save_for_backward(x.shape, y.shape)
- return ctx.op.binary_op(BinaryOps.SUB, x, y, ctx.buffer(binary_broadcast(x.shape, y.shape)))
+ return ctx.binary_op(BinaryOps.SUB, x, y, ctx.buffer(binary_broadcast(x.shape, y.shape)))
def backward(ctx, grad_output):
shape_x, shape_y = ctx.saved_tensors
- neg_grad_output = ctx.op.unary_op(UnaryOps.NEG, grad_output, ctx.buffer(grad_output.shape))
+ neg_grad_output = ctx.unary_op(UnaryOps.NEG, grad_output, ctx.buffer(grad_output.shape))
return unbroadcast(ctx, grad_output, shape_x) if ctx.needs_input_grad[0] else None, \
unbroadcast(ctx, 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 ctx.op.binary_op(BinaryOps.MUL, x, y, ctx.buffer(binary_broadcast(x.shape, y.shape)))
+ return ctx.binary_op(BinaryOps.MUL, x, y, ctx.buffer(binary_broadcast(x.shape, y.shape)))
def backward(ctx, grad_output):
x,y = ctx.saved_tensors
tmp = ctx.buffer(grad_output.shape)
- grad_x = unbroadcast(ctx, ctx.op.binary_op(BinaryOps.MUL, y, grad_output, tmp), x.shape) if ctx.needs_input_grad[0] else None
- grad_y = unbroadcast(ctx, ctx.op.binary_op(BinaryOps.MUL, x, grad_output, tmp), y.shape) if ctx.needs_input_grad[1] else None
+ grad_x = unbroadcast(ctx, ctx.binary_op(BinaryOps.MUL, y, grad_output, tmp), x.shape) if ctx.needs_input_grad[0] else None
+ grad_y = unbroadcast(ctx, ctx.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 = ctx.buffer(binary_broadcast(x.shape, y.shape))
ctx.save_for_backward(x, y, ret)
- return ctx.op.binary_op(BinaryOps.POW, x, y, ret)
+ return ctx.binary_op(BinaryOps.POW, x, y, ret)
def backward(ctx, grad_output):
x,y,powxy = ctx.saved_tensors
tmp = ctx.buffer(grad_output.shape)
- ctx.op.binary_op(BinaryOps.DIV, x, powxy, tmp) # pow(x,y)/x
- ctx.op.binary_op(BinaryOps.MUL, y, tmp, tmp) # y * pow(x,y)/x
- grad_x = unbroadcast(ctx, ctx.op.binary_op(BinaryOps.MUL, grad_output, tmp, tmp), x.shape) if ctx.needs_input_grad[0] else None
- log_x = ctx.op.unary_op(UnaryOps.LOG, x, ctx.buffer(x.shape))
- ctx.op.binary_op(BinaryOps.MUL, log_x, powxy, tmp) # log(x) * pow(x,y)
- grad_y = unbroadcast(ctx, ctx.op.binary_op(BinaryOps.MUL, grad_output, tmp, tmp), y.shape) if ctx.needs_input_grad[1] else None
+ ctx.binary_op(BinaryOps.DIV, x, powxy, tmp) # pow(x,y)/x
+ ctx.binary_op(BinaryOps.MUL, y, tmp, tmp) # y * pow(x,y)/x
+ grad_x = unbroadcast(ctx, ctx.binary_op(BinaryOps.MUL, grad_output, tmp, tmp), x.shape) if ctx.needs_input_grad[0] else None
+ log_x = ctx.unary_op(UnaryOps.LOG, x, ctx.buffer(x.shape))
+ ctx.binary_op(BinaryOps.MUL, log_x, powxy, tmp) # log(x) * pow(x,y)
+ grad_y = unbroadcast(ctx, ctx.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 *************
@@ -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.movement_op(MovementOps.RESHAPE, x, ctx.buffer(shape))
+ return ctx.movement_op(MovementOps.RESHAPE, x, ctx.buffer(shape))
def backward(ctx, grad_output):
in_shape, = ctx.saved_tensors
- return ctx.op.movement_op(MovementOps.RESHAPE, grad_output, ctx.buffer(in_shape))
+ return ctx.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.movement_op(MovementOps.PERMUTE, x, ret, order)
+ return ctx.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.movement_op(MovementOps.PERMUTE, grad_output, ret, norder)
+ return ctx.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.movement_op(MovementOps.SLICE, x, ret, arg)
+ return ctx.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.movement_op(MovementOps.SLICE, grad_output, ret, narg)
+ return ctx.movement_op(MovementOps.SLICE, grad_output, ret, narg)
# ************* processing ops *************
@@ -162,10 +162,10 @@ class Conv2D(Function):
def forward(ctx, x, w, stride=1, groups=1):
C = get_conv_args(x.shape, w.shape, stride, groups)
ctx.save_for_backward(x,w,(C.ys,C.xs), C.groups)
- return ctx.op.processing_op(ProcessingOps.CONV, x, w, ctx.buffer((C.bs, C.groups*C.rcout, C.oy, C.ox)), (C.ys,C.xs), C.groups)
+ return ctx.processing_op(ProcessingOps.CONV, x, w, ctx.buffer((C.bs, C.groups*C.rcout, C.oy, C.ox)), (C.ys,C.xs), C.groups)
def backward(ctx, grad_output):
x, w, stride, groups = ctx.saved_tensors
- dx = ctx.op.processing_op(ProcessingOps.CONVT, grad_output, w, ctx.buffer(x.shape), stride, groups) if ctx.needs_input_grad[0] else None
- dw = ctx.op.processing_op(ProcessingOps.CONVDW, x, grad_output, ctx.buffer(w.shape), stride, groups) if ctx.needs_input_grad[1] else None
+ dx = ctx.processing_op(ProcessingOps.CONVT, grad_output, w, ctx.buffer(x.shape), stride, groups) if ctx.needs_input_grad[0] else None
+ dw = ctx.processing_op(ProcessingOps.CONVDW, x, grad_output, ctx.buffer(w.shape), stride, groups) if ctx.needs_input_grad[1] else None
return dx, dw
\ No newline at end of file

ADD
tinygrad/ops.py
+# TODO: move Device to here and proxy buffer call
+import os
+DEBUG = int(os.getenv("PRINT_LLOPS", "0"))
+class Ops:
+ def unary_op(ctx, op, x, ret):
+ if DEBUG: print(f"{op} : {x.shape} -> {ret.shape}")
+ return ctx.op.unary_op(op, x, ret)
+
+ def reduce_op(ctx, op, inp, ret):
+ if DEBUG: print(f"{op} : {inp.shape} -> {ret.shape}")
+ return ctx.op.reduce_op(op, inp, ret)
+
+ def binary_op(ctx, op, x, y, ret):
+ if DEBUG: print(f"{op} : {x.shape} + {y.shape} -> {ret.shape}")
+ return ctx.op.binary_op(op, x, y, ret)
+
+ def movement_op(ctx, op, inp, ret, arg=None):
+ if DEBUG: print(f"{op} : {inp.shape} -> {ret.shape}")
+ return ctx.op.movement_op(op, inp, ret, arg)
+
+ def processing_op(ctx, op, x, y, ret, stride, groups):
+ if DEBUG: print(f"{op} : {x.shape} + {y.shape} -> {ret.shape}")
+ return ctx.op.processing_op(op, x, y, ret, stride, groups)
\ No newline at end of file

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
@@ -376,7 +376,8 @@ class Tensor:
return y.div((y*y).mean(axis=-1, keepdim=True).add(eps).sqrt())
# An instantiation of the Function is the Context
-class Function:
+from tinygrad.ops import Ops
+class Function(Ops):
def __new__(cls, *args, **kwargs):
cls.forward = staticmethod(cls.forward)
cls.backward = staticmethod(cls.backward)

Ops layer of indirection