From 498b4d2f27090089556ef81eb03f1b95eaad959e Mon Sep 17 00:00:00 2001 From: George Hotz Date: Tue, 10 Nov 2020 15:35:30 -0800 Subject: [PATCH] i32 and reduce line count a bit --- tinygrad/opsgpu.py | 65 +++++++++++++++++++++------------------------- 1 file changed, 29 insertions(+), 36 deletions(-) diff --git a/tinygrad/opsgpu.py b/tinygrad/opsgpu.py index 84a35724bb..675c4c6d24 100644 --- a/tinygrad/opsgpu.py +++ b/tinygrad/opsgpu.py @@ -26,6 +26,8 @@ def clbuild(cl_ctx, prg): def uint2(x, y): return np.array((x,y), dtype=cl.cltypes.uint2) +def i32(x): + return np.int32(x) def cl_subsample_krnl_build(cl_ctx, iter_op, result_op, init_val=0): prg = """ @@ -54,7 +56,7 @@ def subsample_op(ctx, input, kernel_size, stride, iter_op, result_op, init_val=0 prg = cl_subsample_krnl_build(ctx.cl_ctx, iter_op, result_op, init_val=init_val) prg.subsample(ctx.cl_queue, (N*C, Yout, Xout), None, ret, input, uint2(Xout, Yout), uint2(Xin, Yin), - uint2(*kernel_size[::-1]), uint2(px, py), np.int32(input.size)) + uint2(*kernel_size[::-1]), uint2(px, py), i32(input.size)) ctx.data = np.empty((N, C, Yout, Xout)) # set shape expectation on tensor instance return ret @@ -77,7 +79,7 @@ def supersample_op(ctx, input, out_shape, kernel_size, result_op): ret = buffer_zeros(ctx, out_shape) prg = cl_supsample_krnl_build(ctx.cl_ctx, result_op) prg.supsample(ctx.cl_queue, (N*C, Yout, Xout), None, - ret, input, uint2(Xout, Yout), uint2(Xin, Yin), uint2(px, py), np.int32(input.size)) + ret, input, uint2(Xout, Yout), uint2(Xin, Yin), uint2(px, py), i32(input.size)) ctx.data = np.empty((N, C, Yout, Xout)) # set shape expectation on tensor instance return ret @@ -108,7 +110,7 @@ def binary_op(ctx, code, x, y): float b = b_g[gid/ydiv]; res_g[gid] = """+code+"""; }""") - prg.binop(ctx.cl_queue, [np.prod(ret.shape)], None, x, y, ret, np.int32(xdiv), np.int32(ydiv)) + prg.binop(ctx.cl_queue, [np.prod(ret.shape)], None, x, y, ret, i32(xdiv), i32(ydiv)) return ret def unary_op(ctx, code, x): @@ -134,7 +136,7 @@ def reduce_op(ctx, code, code2, input, osize): } res_g[gid] = """+code2+"""; }""") - prg.reduce(ctx.cl_queue, osize, None, input, np.int32(np.prod(input.shape) // np.prod(osize)), ret) + prg.reduce(ctx.cl_queue, osize, None, input, i32(np.prod(input.shape) // np.prod(osize)), ret) return ret # ***** now for the ops themselves ***** @@ -156,8 +158,7 @@ class Sub(Function): @staticmethod def backward(ctx, grad_output): - not_grad_output = unary_op(ctx, '-a', grad_output) - return grad_output, not_grad_output + return grad_output, unary_op(ctx, '-a', grad_output) register('sub', Sub, gpu=True) class Mul(Function): @@ -214,10 +215,7 @@ class Dot(Function): @staticmethod def forward(ctx, input, weight): assert input.shape[1] == weight.shape[0] - isize = np.int32(input.shape[0]) - msize = np.int32(input.shape[1]) - osize = np.int32(weight.shape[1]) - one = np.int32(1) + isize, msize, osize = i32(input.shape[0]), i32(input.shape[1]), i32(weight.shape[1]) ret = buffer_new(ctx, (isize, osize)) prg = clbuild(ctx.cl_ctx, """ @@ -242,16 +240,13 @@ class Dot(Function): # (isize,msize) x (msize,osize) = (isize,osize) prg.matmul(ctx.cl_queue, [isize, osize], None, input, weight, ret, - msize, one, msize, one, osize, osize) + msize, i32(1), msize, i32(1), osize, osize) return ret @staticmethod def backward(ctx, grad_output): input, weight, prg = ctx.saved_tensors - isize = np.int32(input.shape[0]) - msize = np.int32(input.shape[1]) - osize = np.int32(weight.shape[1]) - one = np.int32(1) + isize, msize, osize = i32(input.shape[0]), i32(input.shape[1]), i32(weight.shape[1]) grad_input = buffer_like(ctx, input) grad_weight = buffer_like(ctx, weight) @@ -259,12 +254,12 @@ class Dot(Function): # (isize,osize) x (msize,osize) = (isize,msize) prg.matmul(ctx.cl_queue, [isize, msize], None, grad_output, weight, grad_input, - osize, one, osize, osize, one, msize) + osize, i32(1), osize, osize, i32(1), msize) # (isize,msize) x (isize,osize) = (msize,osize) prg.matmul(ctx.cl_queue, [msize, osize], None, input, grad_output, grad_weight, - one, msize, isize, one, osize, osize) + i32(1), msize, isize, i32(1), osize, osize) return grad_input, grad_weight register('dot', Dot, gpu=True) @@ -294,8 +289,8 @@ class Pad2D(Function): ctx.save_for_backward(padding, prg) prg.pad2d(ctx.cl_queue, [bs*cin, iy, ix], None, x, ret, - np.int32(0), np.int32(0), np.int32(padding[2]), np.int32(padding[0]), - np.int32(oy), np.int32(ox), np.int32(iy), np.int32(ix) + i32(0), i32(0), i32(padding[2]), i32(padding[0]), + i32(oy), i32(ox), i32(iy), i32(ix) ) return ret @@ -307,8 +302,8 @@ class Pad2D(Function): ret = buffer_new(ctx, (bs, cin, oy, ox)) prg.pad2d(ctx.cl_queue, [bs*cin, oy, ox], None, grad_output, ret, - np.int32(padding[2]), np.int32(padding[0]), np.int32(0), np.int32(0), - np.int32(oy), np.int32(ox), np.int32(iy), np.int32(ix) + i32(padding[2]), i32(padding[0]), i32(0), i32(0), + i32(oy), i32(ox), i32(iy), i32(ix) ) return ret register('pad2d', Pad2D, gpu=True) @@ -368,26 +363,24 @@ register('sigmoid', Sigmoid, gpu=True) class AvgPool2D(Function): @staticmethod def forward(ctx, input, kernel_size=(2, 2)): - iter_op = "group_res += input[iid]" - result_op = "group_res / (kernel_size.x * kernel_size.y)" - ret = subsample_op(ctx, input, kernel_size, kernel_size, iter_op, result_op) + ret = subsample_op(ctx, input, kernel_size, kernel_size, + iter_op="group_res += input[iid]", result_op="group_res / (kernel_size.x * kernel_size.y)") ctx.save_for_backward(input.shape) return ret @staticmethod def backward(ctx, grad_output): orig_shape, = ctx.saved_tensors - result_op = "input[iid] / (kernel_size.x * kernel_size.y)" - return supersample_op(ctx, grad_output, orig_shape, ctx.kernel_size, result_op) + return supersample_op(ctx, grad_output, orig_shape, ctx.kernel_size, + result_op="input[iid] / (kernel_size.x * kernel_size.y)") register('avg_pool2d', AvgPool2D, gpu=True) class MaxPool2D(Function): @staticmethod def forward(ctx, input, kernel_size=(2, 2)): - init_val = "FLT_MIN" - iter_op = "group_res = max(group_res, input[iid])" - result_op = "group_res" - return subsample_op(ctx, input, kernel_size, kernel_size, iter_op, result_op, init_val=init_val) + return subsample_op(ctx, input, kernel_size, kernel_size, + iter_op="group_res = max(group_res, input[iid])", + result_op="group_res", init_val="FLT_MIN") @staticmethod def backward(ctx, grad_output): @@ -424,7 +417,7 @@ class LogSoftmax(Function): grad_input[gidsz + gid2] = grad_output[gidsz + gid2] - exp(output[gidsz + gid2]) * acc; }""") prg.lsmsub2(ctx.cl_queue, [grad_output.shape[0], grad_output.shape[1]], None, - grad_output, output, np.int32(grad_output.shape[1]), grad_input) + grad_output, output, i32(grad_output.shape[1]), grad_input) return grad_input register('logsoftmax', LogSoftmax, gpu=True) @@ -477,11 +470,11 @@ class Conv2D(Function): prg.conv(ctx.cl_queue, [bs*groups*rcout, oy, ox], None, x, w, ret, - np.int32(H), np.int32(W), - np.int32(groups), np.int32(rcout), np.int32(cin), - np.int32(oy), np.int32(ox), - np.int32(iy), np.int32(ix), - np.int32(ys), np.int32(xs) + i32(H), i32(W), + i32(groups), i32(rcout), i32(cin), + i32(oy), i32(ox), + i32(iy), i32(ix), + i32(ys), i32(xs) ) return ret