support both asymmetric and negative padding
MODIFY
test/test_ops.py
test/test_ops.py
@@ -266,6 +266,27 @@ class TestOps(unittest.TestCase):
lambda x,w: torch.nn.functional.conv2d(x,w,stride=stride).relu(),
lambda x,w: Tensor.conv2d(x,w,stride=(2,1)).relu(), atol=1e-4)
+ def test_negative_padding_conv2d(self):
+ n,k = 10, 3
+ helper_test_op([(1,1,n,n), (1,1,k,k)],
+ lambda x,w: torch.nn.functional.conv2d(x[:, :, 1:-1, 1:-1],w).relu(),
+ lambda x,w: Tensor.conv2d(x,w,padding=-1).relu(), atol=1e-4)
+ helper_test_op([(1,1,n,n), (1,1,k,k)],
+ lambda x,w: torch.nn.functional.conv2d(x[:, :, 1:, 1:],w).relu(),
+ lambda x,w: Tensor.conv2d(x,w,padding=(-1,0,-1,0)).relu(), atol=1e-4)
+
+ def test_asymmetric_padding_conv2d(self):
+ for p in [(0,1,0,1), (2,1,2,1), (2,0,2,1)]:
+ with self.subTest(padding := p):
+ for n in [3,4]:
+ for k in [2]:
+ helper_test_op([(1,1,n,n), (1,1,k,k)],
+ lambda x,w: torch.nn.functional.conv2d(torch.nn.functional.pad(x, p),w).relu(),
+ lambda x,w: Tensor.conv2d(x,w,padding=p).relu(), atol=1e-4)
+ helper_test_op([(1,1,n,n), (1,1,k,k)],
+ lambda x,w: torch.nn.functional.conv2d(torch.nn.functional.pad(x, p),w).relu(),
+ lambda x,w: Tensor.conv2d(x,w,padding=p).relu(), atol=1e-4)
+
def test_padded_conv2d(self):
bs = 4
cin = 3

MODIFY
tinygrad/helpers.py
tinygrad/helpers.py
@@ -5,18 +5,19 @@ def prod(x): return math.prod(x)
def reduce_shape(shape, axis):
return [1 if i in axis else shape[i] for i in range(len(shape))]
-ConvArgs = namedtuple('ConvArgs', ['H', 'W', 'groups', 'rcout', 'cin', 'oy', 'ox', 'iy', 'ix', 'ys', 'xs', 'bs', 'cout', 'py', 'px', 'dy', 'dx', 'out_shape'])
+ConvArgs = namedtuple('ConvArgs', ['H', 'W', 'groups', 'rcout', 'cin', 'oy', 'ox', 'iy', 'ix', 'ys', 'xs', 'bs', 'cout', 'py', 'py_', 'px', 'px_', 'dy', 'dx', 'out_shape'])
def get_conv_args(x_shape, w_shape, stride=1, groups=1, padding=0, dilation=1):
# TODO: https://docs.nvidia.com/deeplearning/performance/dl-performance-convolutional/index.html#tensor-layout
cout,cin,H,W = w_shape
ys,xs = (stride, stride) if isinstance(stride, int) else stride
- py,px = (padding, padding) if isinstance(padding, int) else padding
+ if not isinstance(padding, int) and len(padding) == 4: px,px_,py,py_ = padding
+ else: py,px = (padding, padding) if isinstance(padding, int) else padding; py_, px_ = py, px
dy,dx = (dilation, dilation) if isinstance(dilation, int) else dilation
bs,cin_,iy,ix = x_shape
# TODO: should be easy to support asymmetric padding by changing output size
# https://pytorch.org/docs/stable/generated/torch.nn.Conv2d.html describes these sizes well
- oy = (iy + 2*py - dy * (H-1) - 1)//ys + 1
- ox = (ix + 2*px - dx * (W-1) - 1)//xs + 1
+ oy = (iy + py + py_ - dy * (H-1) - 1)//ys + 1
+ ox = (ix + px + px_ - dx * (W-1) - 1)//xs + 1
if cin*groups != cin_: raise Exception(f"Input Tensor shape {x_shape} does not match the shape of the weights {w_shape}. ({cin*groups} vs. {cin_})")
assert cout % groups == 0
- return ConvArgs(H, W, groups, cout//groups, cin, oy, ox, iy, ix, ys, xs, bs, cout, py, px, dy, dx, (bs, cout, oy, ox))
+ return ConvArgs(H, W, groups, cout//groups, cin, oy, ox, iy, ix, ys, xs, bs, cout, py, py_, px, px_, dy, dx, (bs, cout, oy, ox))

MODIFY
tinygrad/llops/ops_cpu.py
tinygrad/llops/ops_cpu.py
@@ -61,7 +61,8 @@ class CPUBuffer(np.ndarray):
def processing_op(x,op,w,C):
assert op == ProcessingOps.CONV, f"{op} isn't supported"
- if C.px > 0 or C.py > 0: x = np.pad(x, [(0,0), (0,0), (C.py, C.py), (C.px, C.px)])
+ if C.px != 0 or C.py != 0 or C.px_ != 0 or C.py_ != 0:
+ x = x.movement_op(MovementOps.SLICE, ((0, x.shape[0]), (0, x.shape[1]), (-C.py, x.shape[2]+C.py_), (-C.px, x.shape[3]+C.px_)))
gx = x.reshape(C.bs,C.groups,C.cin,x.shape[2],x.shape[3])
tx = np.lib.stride_tricks.as_strided(gx,
shape=(C.bs, C.groups, C.cin, C.oy, C.ox, C.H, C.W),

MODIFY
tinygrad/llops/ops_gpu.py
tinygrad/llops/ops_gpu.py
@@ -126,7 +126,7 @@ class GPUBuffer:
if C is not None:
ints = ''.join(f"int {x} = {getattr(C, x)};" for x in ["H", "W", "ys", "xs", "dx", "dy", "px", "py", "groups", "rcout", "cin"])
params = [(f"int {x}", getattr(C, x)) for x in ["oy", "ox", "iy", "ix"]]
- if C.px == 0 and C.py == 0: options.append("-DALLVALID")
+ if C.px == 0 and C.py == 0 and C.px_ == 0 and C.py_ == 0: options.append("-DALLVALID")
if C.oy == 1 and C.ox == 1: options.append("-DONEBYONE")
global_size = [C.bs*C.cout, C.oy, C.ox]
assert bufs[0][0] == "input" and bufs[1][0] == "weight"

MODIFY
tinygrad/llops/ops_torch.py
tinygrad/llops/ops_torch.py
import torch
import numpy as np
from tinygrad.llops.ops_cpu import CPUBuffer
-from tinygrad.ops import ProcessingOps
+from tinygrad.ops import MovementOps, ProcessingOps
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
class TorchBuffer(torch.Tensor):
@@ -23,4 +23,6 @@ class TorchBuffer(torch.Tensor):
def processing_op(x,op,w,C):
assert op == ProcessingOps.CONV, f"{op} isn't supported"
- return torch.conv2d(x, w, stride=(C.ys, C.xs), groups=C.groups, dilation=(C.dy, C.dx), padding=(C.py, C.px))
+ if C.px != C.px_ or C.py != C.py_: padding, x = 0, x.movement_op(MovementOps.SLICE, ((0, x.shape[0]), (0, x.shape[1]), (-C.py, x.shape[2]+C.py_), (-C.px, x.shape[3]+C.px_)))
+ else: padding = (C.py, C.px)
+ return torch.conv2d(x, w, stride=(C.ys, C.xs), groups=C.groups, dilation=(C.dy, C.dx), padding=padding)

MODIFY
tinygrad/mlops.py
tinygrad/mlops.py
@@ -203,8 +203,8 @@ class Conv2D(Function):
wt = ctx.movement_op(MovementOps.FLIP, wt, (3, 4))
wt = ctx.movement_op(MovementOps.PERMUTE, wt, (0, 2, 1, 3, 4))
wt = ctx.movement_op(MovementOps.RESHAPE, wt, (C.groups*C.cin, C.rcout, C.H, C.W))
- Cdx = get_conv_args(xt.shape, wt.shape, dilation=(C.dy, C.dx), padding=((C.H-1)*C.dy-C.py,(C.W-1)*C.dx-C.px), groups=C.groups)
- # TODO: this shape can be wrong. support asymmetric padding to remove the slice
+ Cdx = get_conv_args(xt.shape, wt.shape, dilation=(C.dy, C.dx), padding=((C.W-1)*C.dx-C.px, (C.W-1)*C.dx-C.px_, (C.H-1)*C.dy-C.py, (C.H-1)*C.dy-C.py_), groups=C.groups)
+ # TODO: this shape can be wrong strided. support asymmetric padding to remove the slice
dx = ctx._conv(xt, wt, Cdx)
dx = ctx.movement_op(MovementOps.SLICE, dx, [(0,s) for s in x.shape])
@@ -215,7 +215,7 @@ class Conv2D(Function):
xdw = ctx.movement_op(MovementOps.RESHAPE, xdw, (C.cin, C.groups*C.bs, C.iy, C.ix))
grad_output_dw = ctx.movement_op(MovementOps.PERMUTE, grad_output, (1,0,2,3))
grad_output_dw = ctx.movement_op(MovementOps.RESHAPE, grad_output_dw, (C.cout, C.bs, C.oy, C.ox))
- Cdw = get_conv_args(xdw.shape, grad_output_dw.shape, padding=(C.py, C.px), stride=(C.dy, C.dx), dilation=(C.ys, C.xs), groups=C.groups)
+ Cdw = get_conv_args(xdw.shape, grad_output_dw.shape, padding=(C.px, C.px_, C.py, C.py_), stride=(C.dy, C.dx), dilation=(C.ys, C.xs), groups=C.groups)
grad_weight = ctx._conv(xdw, grad_output_dw, Cdw)
grad_weight = ctx.movement_op(MovementOps.PERMUTE, grad_weight, (1,0,2,3))
# TODO: remove this slice using asymmetric padding

support both asymmetric and negative padding