yayay test_sgd_gpu passes
MODIFY
test/test_ops.py
test/test_ops.py
@@ -58,7 +58,7 @@ class TestOps(unittest.TestCase):
def test_dot(self):
helper_test_op([(45,65), (65,100)], lambda x,y: x.matmul(y), Tensor.dot, atol=1e-5, gpu=self.gpu)
def test_sum(self):
- helper_test_op([(45,1)], lambda x: x.sum(), Tensor.sum, atol=1e-5, gpu=self.gpu)
+ helper_test_op([(45,3)], lambda x: x.sum(), Tensor.sum, atol=1e-4, gpu=self.gpu)
def test_logsoftmax(self):
helper_test_op([(45,65)], lambda x: torch.nn.LogSoftmax(dim=1)(x), Tensor.logsoftmax, atol=1e-5, gpu=self.gpu)

MODIFY
tinygrad/opsgpu.py
tinygrad/opsgpu.py
@@ -18,10 +18,6 @@ def buffer_like(ctx, x):
def clbuild(cl_ctx, prg):
return cl.Program(cl_ctx, prg).build()
-@functools.lru_cache
-def clreduce(cl_ctx, *args, **kwargs):
- return ReductionKernel(cl_ctx, *args, **kwargs)
-
def binary_op(ctx, code, x, y):
ret = buffer_like(ctx, x)
prg = clbuild(ctx.cl_ctx, """
@@ -111,11 +107,20 @@ class Sum(Function):
@staticmethod
def forward(ctx, input):
ctx.save_for_backward(input)
- krnl = clreduce(ctx.cl_ctx, np.float32, neutral="0", reduce_expr="a+b",
- map_expr="x[i]", arguments="__global float *x")
- ret = krnl(pycl_array.Array(ctx.cl_queue, input.size, dtype=input.dtype, data=input)).data
- ret.shape = (1,)
- ret.dtype = np.float32
+
+ ret = buffer_new(ctx, (1,))
+ prg = clbuild(ctx.cl_ctx, """
+ __kernel void sum(
+ __global const float *a_g, int sz, __global float *res_g)
+ {
+ float out = 0.0;
+ for (int x = 0; x < sz; x++) {
+ out += a_g[x];
+ }
+ res_g[0] = out;
+ }
+ """)
+ prg.sum(ctx.cl_queue, [input.shape[0]], None, input, np.int32(np.prod(input.shape)), ret)
return ret
@staticmethod
@@ -184,7 +189,7 @@ class Dot(Function):
# (isize,msize) x (isize,osize) = (msize,osize)
prg.matmul(ctx.cl_queue, [msize, osize], None,
input, grad_output, grad_weight,
- one, msize, isize, one, isize, osize)
+ one, msize, isize, one, osize, osize)
return grad_input, grad_weight
register('dot', Dot, gpu=True)

MODIFY
tinygrad/tensor.py
tinygrad/tensor.py
@@ -106,7 +106,10 @@ class Tensor:
if self.gpu:
data = np.empty(self.shape, dtype=np.float32)
cl.enqueue_copy(cl_queue, data, self.data)
- return Tensor(data)
+ ret = Tensor(data)
+ if self.grad:
+ ret.grad = self.grad.cpu()
+ return ret
else:
return self
@@ -123,7 +126,10 @@ class Tensor:
data = cl.Buffer(cl_ctx, cl.mem_flags.READ_ONLY | cl.mem_flags.COPY_HOST_PTR, hostbuf=self.data)
data.shape = self.shape
data.dtype = self.data.dtype
- return Tensor(data)
+ ret = Tensor(data)
+ if self.grad:
+ ret.grad = self.grad.cuda()
+ return ret
else:
return self

yayay test_sgd_gpu passes