gpu speedup, tests work on M1
MODIFY
test/test_net_speed.py
test/test_net_speed.py
@@ -40,7 +40,6 @@ class TestConvSpeed(unittest.TestCase):
mp = torch.nn.MaxPool2d((2,2))
lsm = torch.nn.LogSoftmax(dim=1)
- with torch.autograd.profiler.profile(record_shapes=True) as tprof:
cnt = 5
fpt, bpt = 0.0, 0.0
for i in range(cnt):
@@ -62,8 +61,6 @@ class TestConvSpeed(unittest.TestCase):
print("torch forward pass: %.3f ms" % fpt_baseline)
print("torch backward pass: %.3f ms" % bpt_baseline)
- print(tprof.key_averages().table(sort_by="self_cpu_time_total", row_limit=10))
-
# ****** tinygrad compare *******
c1 = Tensor(c1.detach().numpy())

MODIFY
test/test_nn.py
test/test_nn.py
@@ -17,6 +17,7 @@ class TestNN(unittest.TestCase):
bn.running_var.data[bn.running_var.data < 0] = 0
# create in torch
+ with torch.no_grad():
tbn = torch.nn.BatchNorm2d(sz).eval()
tbn.weight[:] = torch.tensor(bn.weight.data)
tbn.bias[:] = torch.tensor(bn.bias.data)

MODIFY
tinygrad/opsgpu.py
tinygrad/opsgpu.py
@@ -4,13 +4,13 @@ import pyopencl as cl
import functools
def buffer_new(ctx, shape):
- res_g = cl.Buffer(ctx.cl_ctx, cl.mem_flags.WRITE_ONLY, 4*np.prod(shape))
+ res_g = cl.Buffer(ctx.cl_ctx, cl.mem_flags.READ_WRITE, 4*np.prod(shape))
res_g.shape = tuple(shape)
res_g.dtype = np.float32
return res_g
def buffer_np(ctx, np_array):
- res_g = cl.Buffer(ctx.cl_ctx, cl.mem_flags.WRITE_ONLY | cl.mem_flags.COPY_HOST_PTR, hostbuf=np_array)
+ res_g = cl.Buffer(ctx.cl_ctx, cl.mem_flags.READ_WRITE | cl.mem_flags.COPY_HOST_PTR, hostbuf=np_array)
res_g.shape = np_array.shape
res_g.dtype = np_array.dtype
return res_g
@@ -93,7 +93,7 @@ def binary_op(ctx, code, x, y):
if not np.all((shape_x == 1) | (shape_y == 1) | (shape_x == shape_y)):
raise Exception(f"binary op unbroadcastable shape mismatch: {x.shape} vs {y.shape}")
shape_ret = np.maximum(shape_x, shape_y)
- ret = buffer_zeros(ctx, shape_ret)
+ ret = buffer_new(ctx, shape_ret)
binop = clbuild(ctx.cl_ctx, "binop", """
__kernel void binop(__global const float *a_g, __global const float *b_g, __global float *res_g, int n_dims, int prod,

gpu speedup, tests work on M1