github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-09 23:18:04 -05:00
Code Issues Packages Projects Releases Wiki Activity

uint2 saves lines

Browse Source
This commit is contained in:
George Hotz
2020-11-09 23:11:20 -08:00
parent 8b23033fa9
commit d41ad2bf37
1 changed files with 6 additions and 9 deletions
Download Patch File Download Diff File Expand all files Collapse all files

15
tinygrad/opsgpu.py
View File

@@ -24,6 +24,9 @@ def buffer_like(ctx, x):
def clbuild(cl_ctx, prg):
return cl.Program(cl_ctx, prg).build()
def uint2(x, y):
return np.array((x,y), dtype=cl.cltypes.uint2)
@functools.lru_cache
def cl_subsample_krnl_build(cl_ctx, iter_op, result_op, init_val=0):
prg = """
@@ -52,13 +55,10 @@ def subsample_op(ctx, input, kernel_size, stride, iter_op, result_op, init_val=0
N, C, Yin, Xin = input.shape
Yout, Xout = (Yin-kernel_size[0])//py+1, (Xin-kernel_size[1])//px+1
ret = buffer_new(ctx, (N, C, Yout, Xout))
osize = np.array((Xout, Yout), dtype=cl.cltypes.uint2)
isize = np.array((Xin, Yin), dtype=cl.cltypes.uint2)
ksize = np.array(kernel_size[::-1], dtype=cl.cltypes.uint2)
strd = np.array((px, py), dtype=cl.cltypes.uint2)
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, osize, isize, ksize, strd, np.int32(input.size))
ret, input, uint2(Xout, Yout), uint2(Xin, Yin),
uint2(*kernel_size[::-1]), uint2(px, py), np.int32(input.size))
ctx.data = np.empty((N, C, Yout, Xout)) # set shape expectation on tensor instance
return ret
@@ -82,12 +82,9 @@ def supersample_op(ctx, input, out_shape, kernel_size, result_op):
(N, C, Yin, Xin), (Yout, Xout) = input.shape, out_shape[2:]
py,px = kernel_size
ret = buffer_zeros(ctx, out_shape)
osize = np.array((Xout, Yout), dtype=cl.cltypes.uint2)
isize = np.array((Xin, Yin), dtype=cl.cltypes.uint2)
ksize = np.array((px, py), dtype=cl.cltypes.uint2)
prg = cl_supsample_krnl_build(ctx.cl_ctx, result_op)
prg.supsample(ctx.cl_queue, (N*C, Yout, Xout), None,
ret, input, osize, isize, ksize, np.int32(input.size))
ret, input, uint2(Xout, Yout), uint2(Xin, Yin), uint2(px, py), np.int32(input.size))
ctx.data = np.empty((N, C, Yout, Xout)) # set shape expectation on tensor instance
return ret