Add conv2d operator

Author: voltjia

src/ntops/kernels/conv2d.py

@@ -0,0 +1,111 @@
+import functools
+
+import ninetoothed.language as ntl
+from ninetoothed import Symbol, Tensor
+
+import ntops.kernels.mm as mm
+
+STRIDE_H = Symbol("stride_h", constexpr=True)
+STRIDE_W = Symbol("stride_w", constexpr=True)
+DILATION_H = Symbol("dilation_h", constexpr=True)
+DILATION_W = Symbol("dilation_w", constexpr=True)
+
+
+def arrangement(
+    input,
+    weight,
+    bias,
+    output,
+    stride_h=None,
+    stride_w=None,
+    dilation_h=None,
+    dilation_w=None,
+    block_size_m=None,
+    block_size_n=None,
+    block_size_k=None,
+):
+    if stride_h is None:
+        stride_h = STRIDE_H
+
+    if stride_w is None:
+        stride_w = STRIDE_W
+
+    if dilation_h is None:
+        dilation_h = DILATION_H
+
+    if dilation_w is None:
+        dilation_w = DILATION_W
+
+    if block_size_m is None:
+        block_size_m = mm.BLOCK_SIZE_M
+
+    if block_size_n is None:
+        block_size_n = mm.BLOCK_SIZE_N
+
+    if block_size_k is None:
+        block_size_k = mm.BLOCK_SIZE_K
+
+    mm_arrangement = functools.partial(
+        mm.arrangement,
+        block_size_m=block_size_m,
+        block_size_n=block_size_n,
+        block_size_k=block_size_k,
+    )
+
+    input_arranged = input.tile(
+        (1, *weight.shape[1:]),
+        strides=(-1, -1, stride_h, stride_w),
+        dilation=(1, 1, dilation_h, dilation_w),
+        floor_mode=True,
+    )
+    input_arranged = input_arranged.squeeze(1)
+    input_arranged.dtype = input_arranged.dtype.squeeze(0)
+    input_arranged = input_arranged.ravel()
+    input_arranged = input_arranged.flatten(end_dim=3).flatten(start_dim=1)
+
+    weight_arranged = weight.flatten(start_dim=1)
+    weight_arranged = weight_arranged.permute((1, 0))
+
+    bias_arranged = bias.permute((0, 2, 3, 1)).flatten(end_dim=3)
+
+    _, _, bias_arranged = mm_arrangement(input_arranged, weight_arranged, bias_arranged)
+
+    output_arranged = output.permute((0, 2, 3, 1)).flatten(end_dim=3)
+
+    input_arranged, weight_arranged, output_arranged = mm_arrangement(
+        input_arranged, weight_arranged, output_arranged
+    )
+
+    return input_arranged, weight_arranged, bias_arranged, output_arranged
+
+
+def application(input, weight, bias, output):
+    mm_output = ntl.zeros(output.shape, dtype=ntl.float32)
+    mm.application(input, weight, mm_output)
+    output = mm_output + bias
+
+
+def premake(
+    stride_h=None,
+    stride_w=None,
+    dilation_h=None,
+    dilation_w=None,
+    dtype=None,
+    block_size_m=None,
+    block_size_n=None,
+    block_size_k=None,
+):
+    arrangement_ = functools.partial(
+        arrangement,
+        stride_h=stride_h,
+        stride_w=stride_w,
+        dilation_h=dilation_h,
+        dilation_w=dilation_w,
+        block_size_m=block_size_m,
+        block_size_n=block_size_n,
+        block_size_k=block_size_k,
+    )
+
+    tensors = tuple(Tensor(4, dtype=dtype) for _ in range(4))
+
+    return arrangement_, application, tensors

src/ntops/torch.py

@@ -13,6 +13,7 @@
 import ntops.kernels.bitwise_or
 import ntops.kernels.bmm
 import ntops.kernels.clamp
+import ntops.kernels.conv2d
 import ntops.kernels.cos
 import ntops.kernels.div
 import ntops.kernels.dropout
@@ -140,6 +141,54 @@ def clamp(input, min=None, max=None, *, out=None):
     return out
 
 
+def conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):
+    if isinstance(stride, int):
+        stride = (stride, stride)
+
+    # TODO: Support `padding != 0`.
+    assert padding == 0, "`padding != 0` is not supported yet."
+
+    if isinstance(padding, str):
+        if padding == "valid":
+            padding = 0
+
+    if isinstance(padding, int):
+        padding = (padding, padding)
+
+    if isinstance(dilation, int):
+        dilation = (dilation, dilation)
+
+    # TODO: Support `groups != 1`.
+    assert groups == 1, "`groups != 1` is not supported yet."
+
+    n, _, h, w = input.shape
+    k, _, r, s = weight.shape
+    p = math.floor((h + 2 * padding[0] - dilation[0] * (r - 1) - 1) / stride[0] + 1)
+    q = math.floor((w + 2 * padding[1] - dilation[1] * (s - 1) - 1) / stride[1] + 1)
+
+    output = torch.empty((n, k, p, q), dtype=input.dtype, device=input.device)
+
+    if bias is None:
+        bias = torch.zeros((k,), dtype=output.dtype, device=output.device)
+
+    bias = bias[None, :, None, None].expand_as(output)
+
+    kernel = _cached_make(ntops.kernels.conv2d.premake)
+
+    kernel(
+        input,
+        weight,
+        bias,
+        output,
+        stride_h=stride[0],
+        stride_w=stride[1],
+        dilation_h=dilation[0],
+        dilation_w=dilation[1],
+    )
+
+    return output
+
+
 def cos(input, *, out=None):
     if out is None:
         out = torch.empty_like(input)

tests/test_conv2d.py

@@ -0,0 +1,31 @@
+import pytest
+import torch
+import torch.nn.functional as F
+
+import ntops.torch
+from tests.skippers import skip_if_cuda_not_available
+
+
+@skip_if_cuda_not_available
+@pytest.mark.parametrize(
+    "dtype, atol, rtol", ((torch.float16, 0.025, 0.025), (torch.float32, 0.01, 0.01))
+)
+@pytest.mark.parametrize("dilation", (1, 2, (2, 3)))
+@pytest.mark.parametrize("stride", (1, 2, (2, 3)))
+@pytest.mark.parametrize("r, s", ((1, 1), (3, 3)))
+@pytest.mark.parametrize("n, c, h, w, k", ((2, 3, 112, 112, 4),))
+def test_cuda(n, c, h, w, k, r, s, stride, dilation, dtype, atol, rtol):
+    device = "cuda"
+
+    input = torch.randn((n, c, h, w), dtype=dtype, device=device)
+    weight = torch.randn((k, c, r, s), dtype=dtype, device=device)
+    bias = torch.randn((k,), dtype=dtype, device=device)
+
+    ninetoothed_output = ntops.torch.conv2d(
+        input, weight, bias=bias, stride=stride, dilation=dilation
+    )
+    reference_output = F.conv2d(
+        input, weight, bias=bias, stride=stride, dilation=dilation
+    )
+
+    assert torch.allclose(ninetoothed_output, reference_output, atol=atol, rtol=rtol)