all buffer alloc in top level
MODIFY
tinygrad/llops/gpu.py
tinygrad/llops/gpu.py
@@ -38,9 +38,6 @@ class GPUBuffer:
cl.enqueue_copy(cl_queue, data, self.cl, is_blocking=True)
return data
-def buffer_new(shape, zero=False):
- return GPUBuffer(shape, hostbuf=None if not zero else np.zeros(shape, dtype=np.float32))
-
def buffer_np(x):
return cl.Buffer(cl_ctx, cl.mem_flags.READ_WRITE | cl.mem_flags.COPY_HOST_PTR, hostbuf=x)
@@ -87,17 +84,7 @@ def binary_op(code, x, y, ret):
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(code, 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(osize)
- if axis is None:
- ret.shape = (1,)
-
+def reduce_op(code, inp, ret, start="0.0"):
# TODO: this is insanely slow
reduce = clbuild("reduce", """
__kernel void reduce(__global const float *a_g, int sz, __global float *res_g, int prod, int n_dims,
@@ -124,17 +111,15 @@ def reduce_op(code, inp, axis=None, start="0.0"):
}
res_g[gid] = out;
}""")
- reduce([np.prod(osize)], None, inp.cl,
- i32(np.prod(inp.shape)//np.prod(osize)), ret.cl,
- i32(np.prod(osize)), i32(len(osize)),
+ reduce([np.prod(ret.shape)], None, inp.cl,
+ i32(np.prod(inp.shape)//np.prod(ret.shape)), ret.cl,
+ i32(np.prod(ret.shape)), i32(len(ret.shape)),
buffer_np(np.array(inp.shape, dtype=np.int32)),
- buffer_np(np.array(osize, dtype=np.int32)))
+ buffer_np(np.array(ret.shape, dtype=np.int32)))
return ret
-def perm_axis(inp, order):
- osize = np.array(inp.shape)[list(order)]
- ret = buffer_new(osize)
+def perm_axis(inp, order, ret):
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) {
@@ -149,16 +134,14 @@ def perm_axis(inp, order):
}
res_g[gid] = a_g[idx];
}""")
- perm([np.prod(osize)], None, inp.cl, ret.cl, i32(len(osize)),
+ perm([np.prod(inp.shape)], None, inp.cl, ret.cl, i32(len(inp.shape)),
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(x, arg):
+def inner_slice(x, arg, ret):
shift = [y[0] for y in arg]
- oshape = [y[1]-y[0] for y in arg]
- 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,

MODIFY
tinygrad/ops/ops_gpu.py
tinygrad/ops/ops_gpu.py
@@ -2,10 +2,13 @@ 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 GPUBuffer
from ..llops.gpu import unary_op, binary_op, reduce_op, perm_axis, inner_slice
from ..llops.gpu import matmul, conv, convdw, convdx
+def buffer_new(shape, zero=False):
+ return GPUBuffer(shape, hostbuf=None if not zero else np.zeros(shape, dtype=np.float32))
+
# ************* unary ops *************
class UnaryOp(Function):
@@ -31,10 +34,15 @@ class Exp(UnaryOp):
# ************* reduce ops *************
+def reduce_shape(shape, axis):
+ osize = np.array(shape)
+ osize[list(axis)] = 1
+ return osize
+
class Sum(Function):
def forward(ctx, input, axis=None):
ctx.save_for_backward(input.shape)
- return reduce_op("out += a", input, axis=axis)
+ return reduce_op("out += a", input, buffer_new(reduce_shape(input.shape, axis)))
def backward(ctx, grad_output):
shape_input, = ctx.saved_tensors
@@ -44,22 +52,24 @@ class Sum(Function):
class Max(Function):
def forward(ctx, input, axis=None):
- ret = reduce_op("out = max(a,out)", input, axis=axis, start="-INFINITY")
+ osize = np.array(input.shape)
+ osize[list(axis)] = 1
+ ret = reduce_op("out = max(a,out)", input, buffer_new(reduce_shape(input.shape, 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("1.0*(a==b)", input, ret, buffer_new(input.shape))
- div = reduce_op("out += a", ret2, axis=axis, start="1e-10")
+ div = reduce_op("out += a", ret2, buffer_new(reduce_shape(ret2.shape, axis)), start="1e-10")
binary_op("a/b", ret2, div, ret2)
return binary_op('a*b', ret2, grad_output, ret2)
# ************* binary ops *************
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("out += a", out, sum_axis)
+ 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 range(len(out.shape))
+ return reduce_op("out += a", out, buffer_new(reduce_shape(out.shape, sum_axis)))
class Add(Function):
def forward(ctx, x, y):
@@ -120,20 +130,25 @@ class Reshape(Function):
class Transpose(Function):
def forward(ctx, x, order=(1,0)):
ctx.save_for_backward(order)
- return perm_axis(x, order)
+ ret = buffer_new(np.array(x.shape)[list(order)])
+ return perm_axis(x, order, ret)
def backward(ctx, grad_output):
- return perm_axis(grad_output, np.argsort(ctx.order))
+ norder = np.argsort(ctx.order)
+ ret = buffer_new(np.array(grad_output.shape)[list(norder)])
+ return perm_axis(grad_output, norder, ret)
class Slice(Function):
def forward(ctx, x, arg=None):
ctx.save_for_backward(x.shape)
- return inner_slice(x, arg)
+ ret = buffer_new([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)]
- return inner_slice(grad_output, narg)
+ ret = buffer_new([y[1]-y[0] for y in narg])
+ return inner_slice(grad_output, narg, ret)
# ************* processing ops *************

all buffer alloc in top level