refactor to pass conv args into llops
MODIFY
test/test_conv.py
test/test_conv.py
@@ -5,9 +5,11 @@ class TestConv(unittest.TestCase):
def test_simple(self):
x = Tensor.ones(1,12,128,256)
w = Tensor.ones(32,12,3,3)
- ret = x.conv2d(w, padding=(1,1))
+ ret = x.conv2d(w, padding=(1,1)).numpy()
# it's not 108 around the padding
- assert (ret.numpy()[:, :, 1:-1, 1:-1] == 108).all()
+ assert (ret[:, :, 1:-1, 1:-1] == 108).all()
+ assert ret[0,0,0,0] == 48
+ assert ret[0,0,0,1] == 72
if __name__ == '__main__':
unittest.main()
\ No newline at end of file

MODIFY
tinygrad/llops/ops_cpu.py
tinygrad/llops/ops_cpu.py
@@ -65,8 +65,7 @@ def get_tx(x, C):
writeable=False,
)
-def conv(x,w,ret,stride,groups):
- C = get_conv_args(x.shape, w.shape, stride, groups)
+def conv(x,w,ret,C):
tx = get_tx(x, C)
tw = w.reshape(C.groups, C.rcout, C.cin, C.H, C.W)
tmp = np.zeros((C.bs,C.groups,C.oy,C.ox,C.rcout),dtype=x.dtype)
@@ -75,8 +74,7 @@ def conv(x,w,ret,stride,groups):
tmp[:,g] += np.tensordot(tx[:,g], tw[g], ((1,4,5),(1,2,3)))
ret[:] = np.moveaxis(tmp,4,2).reshape(C.bs, C.groups*C.rcout, C.oy, C.ox)
-def convdw(x,grad_output,dw,stride,groups):
- C = get_conv_args(x.shape, dw.shape, stride, groups)
+def convdw(x,grad_output,dw,C):
tx = get_tx(x, C)
ggg = grad_output.reshape(C.bs, C.groups, C.rcout, C.oy, C.ox)
gdw = dw.reshape((C.groups, C.rcout, C.cin, C.H, C.W))
@@ -85,8 +83,7 @@ def convdw(x,grad_output,dw,stride,groups):
#'ikYX,ijYXyx -> kjyx'
gdw[g] += np.tensordot(ggg[:,g], tx[:,g], ((0,2,3),(0,2,3)))
-def convdx(grad_output,w,dx,stride,groups):
- C = get_conv_args(dx.shape, w.shape, stride, groups)
+def convdx(grad_output,w,dx,C):
ggg = grad_output.reshape(C.bs, C.groups, C.rcout, C.oy, C.ox)
tw = w.reshape(C.groups, C.rcout, C.cin, C.H, C.W)
gdx = dx.reshape((C.bs, C.groups, C.cin, C.iy, C.ix))
@@ -99,7 +96,7 @@ def convdx(grad_output,w,dx,stride,groups):
tg = np.dot(ggg[:,g,:,Y,X].reshape(C.bs, -1), tw[g].reshape(C.rcout, -1))
gdx[:, g, :, iY:iY+C.H, iX:iX+C.W] += tg.reshape((C.bs, C.cin, C.H, C.W))
-def processing_op(op,a,b,ret,stride,groups):
- if op == ProcessingOps.CONV: conv(a,b,ret,stride,groups)
- elif op == ProcessingOps.CONVT: convdx(a,b,ret,stride,groups)
- elif op == ProcessingOps.CONVDW: convdw(a,b,ret,stride,groups)
+def processing_op(op,a,b,ret,C):
+ if op == ProcessingOps.CONV: conv(a,b,ret,C)
+ elif op == ProcessingOps.CONVT: convdx(a,b,ret,C)
+ elif op == ProcessingOps.CONVDW: convdw(a,b,ret,C)

MODIFY
tinygrad/llops/ops_gpu.py
tinygrad/llops/ops_gpu.py
@@ -195,8 +195,7 @@ def movement_op(op, x, ret, arg=None):
elif op == MovementOps.PERMUTE: perm_axis(x, arg, ret)
elif op == MovementOps.SLICE: inner_slice(x, arg, ret)
-def conv(x,w,ret,stride,groups):
- C = get_conv_args(x.shape, w.shape, stride, groups)
+def conv(x,w,ret,C):
# input = (bs, groups, cin, iy, ix)
# weight = (groups, rcout, cin, H, W)
# output = (bs, groups, rcout, oy, ox)
@@ -229,8 +228,7 @@ def conv(x,w,ret,stride,groups):
# tensw = (groups*rcout, cin, H, W)
# ggg = (bs, groups*rout, oy, ox)
-def convdw(x,grad_output,dw,stride,groups):
- C = get_conv_args(x.shape, dw.shape, stride, groups)
+def convdw(x,grad_output,dw,C):
convdw_prg = clbuild("convdw", """
__kernel void convdw(__global const float *tensx, __global const float *ggg, __global float *dw,
int H, int W, int groups, int rcout, int cin, int oy, int ox, int iy, int ix, int ys, int xs, int bs) {
@@ -252,8 +250,7 @@ def convdw(x,grad_output,dw,stride,groups):
}""")
convdw_prg([C.groups*C.rcout*C.cin, C.H, C.W], None, x.cl, grad_output.cl, dw.cl, *[i32(x) for x in C])
-def convdx(grad_output,w,dx,stride,groups):
- C = get_conv_args(dx.shape, w.shape, stride, groups)
+def convdx(grad_output,w,dx,C):
convdx_prg = clbuild("convdx", """
__kernel void convdx(__global const float *tensw, __global const float *ggg, __global float *dx,
int H, int W, int groups, int rcout, int cin, int oy, int ox, int iy, int ix, int ys, int xs, int bs) {
@@ -280,7 +277,7 @@ def convdx(grad_output,w,dx,stride,groups):
""")
convdx_prg([C.bs, C.groups, C.cin], None, w.cl, grad_output.cl, dx.cl, *[i32(x) for x in C])
-def processing_op(op,a,b,ret,stride,groups):
- if op == ProcessingOps.CONV: conv(a,b,ret,stride,groups)
- elif op == ProcessingOps.CONVT: convdx(a,b,ret,stride,groups)
- elif op == ProcessingOps.CONVDW: convdw(a,b,ret,stride,groups)
+def processing_op(op,a,b,ret,C):
+ if op == ProcessingOps.CONV: conv(a,b,ret,C)
+ elif op == ProcessingOps.CONVT: convdx(a,b,ret,C)
+ elif op == ProcessingOps.CONVDW: convdw(a,b,ret,C)

MODIFY
tinygrad/llops/ops_torch.py
tinygrad/llops/ops_torch.py
@@ -26,28 +26,27 @@ from tinygrad.llops.ops_cpu import unary_op, binary_op, reduce_op, movement_op
from tinygrad.helpers import get_conv_args
from tinygrad.ops import ProcessingOps
-def convdw(x,grad_output,dw,stride,groups):
+def convdw(x,grad_output,dw,C):
# NOTE: torch.nn.grad.conv2d_weight is wrong for groups in pytorch, wonder who it affects
# https://github.com/pytorch/pytorch/issues/51430
- C = get_conv_args(x.shape, dw.shape, stride, groups)
grad_output = grad_output.reshape(C.bs, C.groups, C.rcout, C.oy, C.ox).repeat(1, 1, C.cin, 1, 1)
grad_output = grad_output.reshape(C.bs * C.groups * C.rcout * C.cin, 1, C.oy, C.ox)
x = x.reshape(1, C.bs * C.groups * C.cin, C.iy, C.ix)
- grad_weight = torch.conv2d(x, grad_output, dilation=stride, groups=C.bs*C.groups*C.cin)
+ grad_weight = torch.conv2d(x, grad_output, dilation=(C.ys, C.xs), groups=C.bs*C.groups*C.cin)
grad_weight = grad_weight.reshape(C.bs, C.groups, C.cin, C.rcout, *grad_weight.shape[2:]).sum(dim=0).transpose(2, 1)
dw[:] = grad_weight.reshape(C.groups*C.rcout, C.cin, *grad_weight.shape[3:])[:, :, :dw.shape[2], :dw.shape[3]]
-def processing_op(op,x,w,ret,stride,groups):
+def processing_op(op,x,w,ret,C):
+ stride, groups = (C.ys, C.xs), C.groups
if op == ProcessingOps.CONV:
ret[:] = torch.conv2d(x, w, stride=stride, groups=groups)
elif op == ProcessingOps.CONVT:
if stride == 1 or stride == (1,1):
# strided needs weird padding: https://github.com/vdumoulin/conv_arithmetic/blob/master/README.md
- C = get_conv_args(ret.shape, w.shape, stride, groups)
w = w.reshape(C.groups, C.rcout, C.cin, C.H, C.W).flip(3, 4).transpose(2, 1).reshape(C.groups*C.cin, C.rcout, C.H, C.W)
ret[:] = torch.conv2d(x, w, padding=(C.H-1,C.W-1), groups=groups)
else:
output_padding = [ret.shape[d+2] - ((x.shape[d+2] - 1) * stride[d] + 1 + (w.shape[d+2] - 1)) for d in range(2)]
ret[:] = torch.conv_transpose2d(x, w, stride=stride, groups=groups, output_padding=output_padding)
elif op == ProcessingOps.CONVDW:
- convdw(x,w,ret,stride,groups)
+ convdw(x,w,ret,C)

MODIFY
tinygrad/mlops.py
tinygrad/mlops.py
@@ -151,10 +151,11 @@ class Conv2D(Function):
def forward(ctx, x, w, stride=1, groups=1):
C = get_conv_args(x.shape, w.shape, stride, groups)
ctx.save_for_backward(x,w,(C.ys,C.xs), C.groups)
- return ctx.processing_op(ProcessingOps.CONV, x, w, (C.bs, C.groups*C.rcout, C.oy, C.ox), (C.ys,C.xs), C.groups)
+ return ctx.processing_op(ProcessingOps.CONV, x, w, (C.bs, C.groups*C.rcout, C.oy, C.ox), C)
def backward(ctx, grad_output):
x, w, stride, groups = ctx.saved_tensors
- dx = ctx.processing_op(ProcessingOps.CONVT, grad_output, w, x.shape, stride, groups) if ctx.needs_input_grad[0] else None
- dw = ctx.processing_op(ProcessingOps.CONVDW, x, grad_output, w.shape, stride, groups) if ctx.needs_input_grad[1] else None
+ C = get_conv_args(x.shape, w.shape, stride, groups)
+ dx = ctx.processing_op(ProcessingOps.CONVT, grad_output, w, x.shape, C) if ctx.needs_input_grad[0] else None
+ dw = ctx.processing_op(ProcessingOps.CONVDW, x, grad_output, w.shape, C) if ctx.needs_input_grad[1] else None
return dx, dw
\ No newline at end of file

MODIFY
tinygrad/ops.py
tinygrad/ops.py
@@ -74,9 +74,9 @@ class Ops:
log_op(op, ret, [x])
return ret
- def processing_op(ctx, op:ProcessingOps, x, y, out_shape, stride, groups):
+ def processing_op(ctx, op:ProcessingOps, x, y, out_shape, C):
# TODO: can we do better than out_shape?
ret = ctx.buffer(out_shape)
- ctx.op.processing_op(op, x, y, ret, stride, groups)
+ ctx.op.processing_op(op, x, y, ret, C)
log_op(op, ret, [x, y])
return ret
\ No newline at end of file

refactor to pass conv args into llops