remove all ctxes
MODIFY
tinygrad/llops/gpu.py
tinygrad/llops/gpu.py
@@ -39,10 +39,10 @@ class GPUBuffer:
cl.enqueue_copy(cl_queue, data, self.cl, is_blocking=True)
return data
-def buffer_new(ctx, shape, zero=False):
+def buffer_new(shape, zero=False):
return GPUBuffer(shape, hostbuf=None if not zero else np.zeros(shape, dtype=np.float32))
-def buffer_np(ctx, x):
+def buffer_np(x):
return cl.Buffer(cl_ctx, cl.mem_flags.READ_WRITE | cl.mem_flags.COPY_HOST_PTR, hostbuf=x)
def clbuffer(hostbuf, shape):
@@ -58,8 +58,8 @@ def clbuild(name, prg):
return run
# x -> ret
-def unary_op(ctx, code, x):
- ret = buffer_new(ctx, x.shape)
+def unary_op(code, x):
+ ret = buffer_new(x.shape)
unop = clbuild("unop", """
__kernel void unop(__global const float *a_g, __global float *res_g) {
int gid = get_global_id(0);
@@ -70,7 +70,7 @@ def unary_op(ctx, code, x):
return ret
@functools.lru_cache
-def get_binop_prg(cl_ctx, code, complist):
+def get_binop_prg(code, complist):
ndims = len(complist)
args = "".join([f", int d{i}" for i in range(ndims)] + [f", int p{i}" for i in range(ndims-1)])
compute_idx_rets = "".join([f"\n int idx_ret{i} = (gid0 / {f'p{i}' if i < ndims-1 else '1'}) % d{i};" for i in range(ndims)])
@@ -87,23 +87,23 @@ def get_binop_prg(cl_ctx, code, complist):
float b = y_g["""+idx_exprs[1]+"""];
res_g[gid0] = """+code+""";\n}""").build()
-def binary_op(ctx, code, x, y):
+def binary_op(code, x, y):
shape_ret, dimlist, complist = binary_broadcast(x.shape, y.shape)
prod_list = np.array(dimlist, dtype=i32)[-1::-1].cumprod(dtype=i32)[-1::-1] # take cumprod from back to front
- prg = get_binop_prg(cl_ctx, code, tuple(complist))
- ret = buffer_new(ctx, shape_ret, zero=True)
+ 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
-def reduce_op(ctx, code, code2, inp, axis=None, start="0.0"):
+def reduce_op(code, code2, inp, axis=None, start="0.0"):
if axis is None:
# full reduce
osize = [1]*len(inp.shape)
else:
osize = np.array(inp.shape)
osize[list(axis)] = 1
- ret = buffer_new(ctx, osize)
+ ret = buffer_new(osize)
if axis is None:
ret.shape = (1,)
@@ -136,14 +136,14 @@ def reduce_op(ctx, code, code2, inp, axis=None, start="0.0"):
reduce([np.prod(osize)], None, inp.cl,
i32(np.prod(inp.shape)//np.prod(osize)), ret.cl,
i32(np.prod(osize)), i32(len(osize)),
- buffer_np(ctx, np.array(inp.shape, dtype=np.int32)),
- buffer_np(ctx, np.array(osize, dtype=np.int32)))
+ buffer_np(np.array(inp.shape, dtype=np.int32)),
+ buffer_np(np.array(osize, dtype=np.int32)))
return ret
-def perm_axis(ctx, inp, order):
+def perm_axis(inp, order):
osize = np.array(inp.shape)[list(order)]
- ret = buffer_new(ctx, osize)
+ ret = buffer_new(osize)
perm = clbuild("perm", """
__kernel void perm(__global const float *a_g, __global float *res_g, int n_axis,
__global const int *shape, __global const int *order) {
@@ -159,15 +159,15 @@ def perm_axis(ctx, inp, order):
res_g[gid] = a_g[idx];
}""")
perm([np.prod(osize)], None, inp.cl, ret.cl, i32(len(osize)),
- buffer_np(ctx, np.array(inp.shape, dtype=np.int32)),
- buffer_np(ctx, np.array(order, dtype=np.int32)))
+ buffer_np(np.array(inp.shape, dtype=np.int32)),
+ buffer_np(np.array(order, dtype=np.int32)))
return ret
# TODO: merge this with perm axis
-def inner_slice(ctx, x, arg):
+def inner_slice(x, arg):
shift = [y[0] for y in arg]
oshape = [y[1]-y[0] for y in arg]
- ret = buffer_new(ctx, oshape)
+ ret = buffer_new(oshape)
gslice = clbuild("gslice", """
__kernel void gslice(__global const float *input, __global float *output, int prod, int n_dims,
__global const int *shape_x, __global const int *shape_ret,
@@ -185,9 +185,9 @@ def inner_slice(ctx, x, arg):
}""")
gslice([np.prod(ret.shape)], None,
x.cl, ret.cl, i32(np.prod(ret.shape)), i32(len(ret.shape)),
- buffer_np(ctx, np.array(x.shape, dtype=np.int32)),
- buffer_np(ctx, np.array(ret.shape, dtype=np.int32)),
- buffer_np(ctx, np.array(shift, dtype=np.int32)))
+ buffer_np(np.array(x.shape, dtype=np.int32)),
+ buffer_np(np.array(ret.shape, dtype=np.int32)),
+ buffer_np(np.array(shift, dtype=np.int32)))
return ret
# c = a@b

MODIFY
tinygrad/ops/ops_gpu.py
tinygrad/ops/ops_gpu.py
import pyopencl as cl
import numpy as np
from ..tensor import Function
-from ..llops.gpu import GPUBuffer, clbuild, buffer_new, unary_op, binary_op, reduce_op, perm_axis, inner_slice, matmul, conv, convdw, convdx
-
-i32 = np.int32
+from ..llops.gpu import GPUBuffer, buffer_new
+from ..llops.gpu import unary_op, binary_op, reduce_op, perm_axis, inner_slice
+from ..llops.gpu import matmul, conv, convdw, convdx
# ************* unary ops *************
class UnaryOp(Function):
def forward(ctx, input):
ctx.save_for_backward(input)
- return unary_op(ctx, ctx.fop, input)
+ return unary_op(ctx.fop, input)
def backward(ctx, grad_output):
input, = ctx.saved_tensors
- return binary_op(ctx, ctx.bop, input, grad_output)
+ return binary_op(ctx.bop, input, grad_output)
class ReLU(UnaryOp):
fop = 'max(a, (float)0.)'
@@ -33,75 +33,75 @@ class Exp(UnaryOp):
class Sum(Function):
def forward(ctx, input, axis=None):
ctx.save_for_backward(input.shape)
- return reduce_op(ctx, "out += a", "out", input, axis=axis)
+ return reduce_op("out += a", "out", input, axis=axis)
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(ctx, 'a', grad_output, buffer_new(ctx, shape_input))
+ return binary_op('a', grad_output, buffer_new(shape_input))
class Max(Function):
def forward(ctx, input, axis=None):
- ret = reduce_op(ctx, "out = max(a,out)", "out", input, axis=axis, start="-INFINITY")
+ ret = reduce_op("out = max(a,out)", "out", input, axis=axis, start="-INFINITY")
ctx.save_for_backward(input, axis, ret)
return ret
def backward(ctx, grad_output):
input, axis, ret = ctx.saved_tensors
- ret2 = binary_op(ctx, "1.0*(a==b)", input, ret)
- div = reduce_op(ctx, "out += a", "out+1e-10", ret2, axis=axis)
- ret3 = binary_op(ctx, "a/b", ret2, div)
- return binary_op(ctx, 'a*b', ret3, grad_output)
+ ret2 = binary_op("1.0*(a==b)", input, ret)
+ div = reduce_op("out += a", "out+1e-10", ret2, axis=axis)
+ ret3 = binary_op("a/b", ret2, div)
+ return binary_op('a*b', ret3, grad_output)
# ************* binary ops *************
-def unbroadcast(ctx, out, in_sh):
+def unbroadcast(out, in_sh):
sum_axis = [i for i in range(len(in_sh)) if in_sh[i]==1 and out.shape[i]>1] if in_sh != (1,) else None
- return reduce_op(ctx, "out += a", "out", out, sum_axis)
+ return reduce_op("out += a", "out", out, sum_axis)
class Add(Function):
def forward(ctx, x, y):
ctx.save_for_backward(x.shape, y.shape)
- return binary_op(ctx, 'a+b', x, y)
+ return binary_op('a+b', x, y)
def backward(ctx, grad_output):
grad_x, grad_y = grad_output, grad_output
shape_x, shape_y = ctx.saved_tensors
- return unbroadcast(ctx, grad_x, shape_x), unbroadcast(ctx, grad_y, shape_y)
+ return unbroadcast(grad_x, shape_x), unbroadcast(grad_y, shape_y)
class Sub(Function):
def forward(ctx, x, y):
ctx.save_for_backward(x.shape, y.shape)
- return binary_op(ctx, 'a-b', x, y)
+ return binary_op('a-b', x, y)
def backward(ctx, grad_output):
- grad_x, grad_y = grad_output, unary_op(ctx, '-a', grad_output)
+ grad_x, grad_y = grad_output, unary_op('-a', grad_output)
shape_x, shape_y = ctx.saved_tensors
- return unbroadcast(ctx, grad_x, shape_x), unbroadcast(ctx, grad_y, shape_y)
+ 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(ctx, 'a*b', x, y)
+ return binary_op('a*b', x, y)
def backward(ctx, grad_output):
x,y = ctx.saved_tensors
- grad_x = binary_op(ctx, 'a*b', y, grad_output)
- grad_y = binary_op(ctx, 'a*b', x, grad_output)
- return unbroadcast(ctx, grad_x, x.shape), unbroadcast(ctx, grad_y, y.shape)
+ grad_x = binary_op('a*b', y, grad_output)
+ grad_y = binary_op('a*b', x, grad_output)
+ 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(ctx, 'pow(a,b)', x, y)
+ return binary_op('pow(a,b)', x, y)
def backward(ctx, grad_output):
x,y = ctx.saved_tensors
- grad_x_inter = binary_op(ctx, 'b * (pow((float)a, (float)(b-1.0)))', x, y)
- grad_x = binary_op(ctx, 'a*b', grad_output, grad_x_inter)
- grad_y_inter = binary_op(ctx, 'pow(a, (float)b) * log(a);', x, y)
- grad_y = binary_op(ctx, 'a*b', grad_output, grad_y_inter)
- return unbroadcast(ctx, grad_x, x.shape), unbroadcast(ctx, grad_y, y.shape)
+ 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)
+ return unbroadcast(grad_x, x.shape), unbroadcast(grad_y, y.shape)
# ************* movement ops *************
@@ -117,38 +117,37 @@ class Reshape(Function):
in_shape, = ctx.saved_tensors
return GPUBuffer(in_shape, hostbuf=grad_output)
-
class Transpose(Function):
def forward(ctx, x, order=(1,0)):
ctx.save_for_backward(order)
- return perm_axis(ctx, x, order)
+ return perm_axis(x, order)
def backward(ctx, grad_output):
- return perm_axis(ctx, grad_output, np.argsort(ctx.order))
+ return perm_axis(grad_output, np.argsort(ctx.order))
class Slice(Function):
def forward(ctx, x, arg=None):
ctx.save_for_backward(x.shape)
- return inner_slice(ctx, x, arg)
+ 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(ctx, grad_output, narg)
+ return inner_slice(grad_output, narg)
# ************* processing ops *************
class Matmul(Function):
def forward(ctx, input, weight):
assert input.shape[-1] == weight.shape[-2]
- ret = buffer_new(ctx, list(input.shape[0:-1])+[weight.shape[-1]])
+ ret = buffer_new(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 = buffer_new(ctx, input.shape)
- grad_weight = buffer_new(ctx, weight.shape)
+ grad_input = buffer_new(input.shape)
+ grad_weight = buffer_new(weight.shape)
matmul(grad_output, weight, grad_input, transpose_b=True)
matmul(input, grad_output, grad_weight, transpose_a=True)
return grad_input, grad_weight
@@ -168,7 +167,7 @@ class Conv2D(Function):
# output buffer
conv_args = H, W, groups, rcout, cin, oy, ox, iy, ix, ys, xs, bs
- return conv(x, w, buffer_new(ctx, (bs, cout, oy, ox)), conv_args)
+ return conv(x, w, buffer_new((bs, cout, oy, ox)), conv_args)
def backward(ctx, grad_output):
bs,_,oy,ox = grad_output.shape
@@ -182,6 +181,6 @@ class Conv2D(Function):
rcout = cout//ctx.groups
conv_args = H, W, ctx.groups, rcout, cin, oy, ox, iy, ix, ys, xs, bs
- dw = convdw(x, grad_output, buffer_new(ctx, (cout, cin, H, W)), conv_args)
- dx = convdx(w, grad_output, buffer_new(ctx, (bs, cin_, iy, ix), zero=True), conv_args)
+ dw = convdw(x, grad_output, buffer_new((cout, cin, H, W)), conv_args)
+ dx = convdx(w, grad_output, buffer_new((bs, cin_, iy, ix), zero=True), conv_args)
return dx, dw

remove all ctxes