factor out unary and binary op allocations
MODIFY
tinygrad/helpers.py
tinygrad/helpers.py
import numpy as np
-def binary_broadcast(x_shape, y_shape):
+def binary_broadcast(x_shape, y_shape, extra=False):
n_dims = max(len(x_shape), len(y_shape))
shape_x, shape_y = np.ones(n_dims, dtype=np.int32), np.ones(n_dims, dtype=np.int32)
shape_x[:len(x_shape)] = np.array(x_shape, dtype=np.int32)
@@ -9,6 +9,7 @@ def binary_broadcast(x_shape, y_shape):
raise Exception(f"binary op unbroadcastable shape mismatch: {x_shape} vs {y_shape}")
shape_ret = np.maximum(shape_x, shape_y)
+ if extra:
dimlist, complist = [], [] # note: len(dimlist) may be less than n_dims
def push(dim, comp):
if len(complist) > 0 and complist[-1] == comp:
@@ -19,4 +20,6 @@ def binary_broadcast(x_shape, y_shape):
push(np.int32(max(shape_x[i], shape_y[i])), (shape_x[i] > 1, shape_y[i] > 1))
return shape_ret, dimlist, complist
+ else:
+ return shape_ret

MODIFY
tinygrad/llops/gpu.py
tinygrad/llops/gpu.py
@@ -52,8 +52,7 @@ def clbuild(name, prg):
return run
# x -> ret
-def unary_op(code, x):
- ret = buffer_new(x.shape)
+def unary_op(code, x, ret):
unop = clbuild("unop", """
__kernel void unop(__global const float *a_g, __global float *res_g) {
int gid = get_global_id(0);
@@ -81,12 +80,11 @@ def get_binop_prg(code, complist):
float b = y_g["""+idx_exprs[1]+"""];
res_g[gid0] = """+code+""";\n}""").build()
-def binary_op(code, x, y):
- shape_ret, dimlist, complist = binary_broadcast(x.shape, y.shape)
+def binary_op(code, x, y, ret):
+ shape_ret, dimlist, complist = binary_broadcast(x.shape, y.shape, True)
+ assert (shape_ret == ret.shape).all()
prod_list = np.array(dimlist, dtype=i32)[-1::-1].cumprod(dtype=i32)[-1::-1] # take cumprod from back to front
-
prg = get_binop_prg(code, tuple(complist))
- ret = buffer_new(shape_ret, zero=True)
prg.binop(cl_queue, [prod_list[0]] if len(dimlist) > 0 else [1], None, x.cl, y.cl, ret.cl, *dimlist, *(prod_list[1:]))
return ret

MODIFY
tinygrad/ops/ops_gpu.py
tinygrad/ops/ops_gpu.py
import pyopencl as cl
import numpy as np
+from tinygrad.helpers import binary_broadcast
from ..tensor import Function
from ..llops.gpu import GPUBuffer, buffer_new
from ..llops.gpu import unary_op, binary_op, reduce_op, perm_axis, inner_slice
@@ -10,11 +11,11 @@ from ..llops.gpu import matmul, conv, convdw, convdx
class UnaryOp(Function):
def forward(ctx, input):
ctx.save_for_backward(input)
- return unary_op(ctx.fop, input)
+ return unary_op(ctx.fop, input, buffer_new(input.shape))
def backward(ctx, grad_output):
input, = ctx.saved_tensors
- return binary_op(ctx.bop, input, grad_output)
+ return binary_op(ctx.bop, input, grad_output, buffer_new(input.shape))
class ReLU(UnaryOp):
fop = 'max(a, (float)0.)'
@@ -38,7 +39,8 @@ class Sum(Function):
def backward(ctx, grad_output):
shape_input, = ctx.saved_tensors
# NOTE: the b buffer_new isn't used, since this is just for broadcast
- return binary_op('a', grad_output, buffer_new(shape_input))
+ ret = buffer_new(shape_input)
+ return binary_op('a', grad_output, ret, ret)
class Max(Function):
def forward(ctx, input, axis=None):
@@ -48,10 +50,10 @@ class Max(Function):
def backward(ctx, grad_output):
input, axis, ret = ctx.saved_tensors
- ret2 = binary_op("1.0*(a==b)", input, ret)
+ ret2 = binary_op("1.0*(a==b)", input, ret, buffer_new(input.shape))
div = reduce_op("out += a", ret2, axis=axis, start="1e-10")
- ret3 = binary_op("a/b", ret2, div)
- return binary_op('a*b', ret3, grad_output)
+ ret3 = binary_op("a/b", ret2, div, buffer_new(input.shape))
+ return binary_op('a*b', ret3, grad_output, buffer_new(input.shape))
# ************* binary ops *************
@@ -62,7 +64,7 @@ def unbroadcast(out, in_sh):
class Add(Function):
def forward(ctx, x, y):
ctx.save_for_backward(x.shape, y.shape)
- return binary_op('a+b', x, y)
+ return binary_op('a+b', x, y, buffer_new(binary_broadcast(x.shape, y.shape)))
def backward(ctx, grad_output):
grad_x, grad_y = grad_output, grad_output
@@ -72,35 +74,35 @@ class Add(Function):
class Sub(Function):
def forward(ctx, x, y):
ctx.save_for_backward(x.shape, y.shape)
- return binary_op('a-b', x, y)
+ return binary_op('a-b', x, y, buffer_new(binary_broadcast(x.shape, y.shape)))
def backward(ctx, grad_output):
- grad_x, grad_y = grad_output, unary_op('-a', grad_output)
+ grad_x, grad_y = grad_output, unary_op('-a', grad_output, buffer_new(grad_output.shape))
shape_x, shape_y = ctx.saved_tensors
return unbroadcast(grad_x, shape_x), unbroadcast(grad_y, shape_y)
class Mul(Function):
def forward(ctx, x, y):
ctx.save_for_backward(x, y)
- return binary_op('a*b', x, y)
+ return binary_op('a*b', x, y, buffer_new(binary_broadcast(x.shape, y.shape)))
def backward(ctx, grad_output):
x,y = ctx.saved_tensors
- grad_x = binary_op('a*b', y, grad_output)
- grad_y = binary_op('a*b', x, grad_output)
+ grad_x = binary_op('a*b', y, grad_output, buffer_new(binary_broadcast(y.shape, grad_output.shape)))
+ grad_y = binary_op('a*b', x, grad_output, buffer_new(binary_broadcast(x.shape, grad_output.shape)))
return unbroadcast(grad_x, x.shape), unbroadcast(grad_y, y.shape)
class Pow(Function):
def forward(ctx, x, y):
ctx.save_for_backward(x, y)
- return binary_op('pow(a,b)', x, y)
+ return binary_op('pow(a,b)', x, y, buffer_new(binary_broadcast(x.shape, y.shape)))
def backward(ctx, grad_output):
x,y = ctx.saved_tensors
- grad_x_inter = binary_op('b * (pow((float)a, (float)(b-1.0)))', x, y)
- grad_x = binary_op('a*b', grad_output, grad_x_inter)
- grad_y_inter = binary_op('pow(a, (float)b) * log(a);', x, y)
- grad_y = binary_op('a*b', grad_output, grad_y_inter)
+ grad_x_inter = binary_op('b * (pow((float)a, (float)(b-1.0)))', x, y, buffer_new(grad_output.shape))
+ grad_x = binary_op('a*b', grad_output, grad_x_inter, buffer_new(grad_output.shape))
+ grad_y_inter = binary_op('pow(a, (float)b) * log(a);', x, y, buffer_new(grad_output.shape))
+ grad_y = binary_op('a*b', grad_output, grad_y_inter, buffer_new(grad_output.shape))
return unbroadcast(grad_x, x.shape), unbroadcast(grad_y, y.shape)
# ************* movement ops *************

factor out unary and binary op allocations