Add conv2d operator

Author: voltjia

src/ntops/kernels/__init__.py

@@ -8,6 +8,7 @@
     bitwise_or,
     bmm,
     clamp,
+    conv2d,
     cos,
     div,
     dropout,
@@ -50,6 +51,7 @@
     "bitwise_or",
     "bmm",
     "clamp",
+    "conv2d",
     "cos",
     "div",
     "dropout",

src/ntops/kernels/conv2d.py

@@ -0,0 +1,144 @@
+import copy
+import functools
+
+import ninetoothed.language as ntl
+from ninetoothed import Symbol, Tensor
+
+from ntops.kernels import mm
+
+
+def arrangement(
+    input,
+    weight,
+    bias,
+    output,
+    input_precision,
+    stride_h=None,
+    stride_w=None,
+    padding_h=None,
+    padding_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 = Symbol("stride_h", constexpr=True)
+
+    if stride_w is None:
+        stride_w = Symbol("stride_w", constexpr=True)
+
+    if padding_h is None:
+        padding_h = Symbol("padding_h", constexpr=True)
+
+    if padding_w is None:
+        padding_w = Symbol("padding_w", constexpr=True)
+
+    if dilation_h is None:
+        dilation_h = Symbol("dilation_h", constexpr=True)
+
+    if dilation_w is None:
+        dilation_w = Symbol("dilation_w", constexpr=True)
+
+    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.pad(
+        ((0, 0), (0, 0), (padding_h, padding_h), (padding_w, padding_w))
+    )
+    input_arranged = input_arranged.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[None, :, None, None].expand(
+        (output.shape[0], -1, output.shape[2], output.shape[3])
+    )
+    bias_arranged = bias_arranged.permute((0, 2, 3, 1)).flatten(end_dim=3)
+
+    output_arranged = output.permute((0, 2, 3, 1)).flatten(end_dim=3)
+
+    _, _, bias_arranged, _ = mm_arrangement(
+        copy.deepcopy(input_arranged),
+        copy.deepcopy(weight_arranged),
+        bias_arranged,
+        copy.deepcopy(input_precision),
+    )
+
+    input_arranged, weight_arranged, output_arranged, input_precision_arranged = (
+        mm_arrangement(
+            input_arranged, weight_arranged, output_arranged, input_precision
+        )
+    )
+
+    return (
+        input_arranged,
+        weight_arranged,
+        bias_arranged,
+        output_arranged,
+        input_precision_arranged,
+    )
+
+
+def application(input, weight, bias, output, input_precision):
+    mm_output = ntl.zeros(output.shape, dtype=ntl.float32)
+    mm.application(input, weight, mm_output, input_precision)
+    output = mm_output + bias
+
+
+def premake(
+    input_precision=None,
+    stride_h=None,
+    stride_w=None,
+    padding_h=None,
+    padding_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,
+        padding_h=padding_h,
+        padding_w=padding_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,
+    )
+
+    input, weight, output = (Tensor(4, dtype=dtype) for _ in range(3))
+    bias = Tensor(1, dtype=dtype)
+    input_precision = Tensor(0, dtype=dtype, constexpr=True, value=input_precision)
+
+    tensors = (input, weight, bias, output, input_precision)
+
+    return arrangement_, application, tensors

src/ntops/torch/__init__.py

@@ -7,6 +7,7 @@
 from ntops.torch.bitwise_or import bitwise_or
 from ntops.torch.bmm import bmm
 from ntops.torch.clamp import clamp
+from ntops.torch.conv2d import conv2d
 from ntops.torch.cos import cos
 from ntops.torch.div import div
 from ntops.torch.dropout import dropout
@@ -49,6 +50,7 @@
     "bitwise_or",
     "bmm",
     "clamp",
+    "conv2d",
     "cos",
     "div",
     "dropout",

src/ntops/torch/conv2d.py

@@ -0,0 +1,51 @@
+import torch
+
+import ntops
+from ntops.torch.pooling import _calculate_output_size
+from ntops.torch.utils import _cached_make, _get_matmul_input_precision
+
+
+def conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1):
+    if isinstance(stride, int):
+        stride = (stride, stride)
+
+    if isinstance(padding, str):
+        if padding == "valid":
+            padding = 0
+
+    if isinstance(padding, int):
+        padding = (padding, padding)
+
+    if isinstance(dilation, int):
+        dilation = (dilation, dilation)
+
+    assert groups == 1, "`groups` is not supported yet."
+
+    n, c, h, w = input.shape
+    k, _, r, s = weight.shape
+
+    p = _calculate_output_size(
+        h, r, stride=stride[0], padding=padding[0], dilation=dilation[0]
+    )
+    q = _calculate_output_size(
+        w, s, stride=stride[1], padding=padding[1], dilation=dilation[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)
+
+    kernel = _cached_make(
+        ntops.kernels.conv2d.premake,
+        stride_h=stride[0],
+        stride_w=stride[1],
+        padding_h=padding[0],
+        padding_w=padding[1],
+        dilation_h=dilation[0],
+        dilation_w=dilation[1],
+    )
+
+    kernel(input, weight, bias, output, _get_matmul_input_precision())
+
+    return output

tests/test_conv2d.py

@@ -0,0 +1,33 @@
+import pytest
+import torch
+import torch.nn.functional as F
+
+import ntops
+from tests.skippers import skip_if_cuda_not_available
+
+
+@skip_if_cuda_not_available
+@pytest.mark.parametrize("device", ("cuda",))
+@pytest.mark.parametrize(
+    "dtype, rtol, atol", ((torch.float32, 1e-5, 1e-5), (torch.float16, 1e-3, 1e-3))
+)
+@pytest.mark.parametrize("dilation", (1, 2, (2, 3)))
+@pytest.mark.parametrize("padding", (0, 1, (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_conv2d(
+    n, c, h, w, k, r, s, stride, padding, dilation, dtype, device, rtol, atol
+):
+    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, padding=padding, dilation=dilation
+    )
+    reference_output = F.conv2d(
+        input, weight, bias=bias, stride=stride, padding=padding, dilation=dilation
+    )
+
+    assert torch.allclose(ninetoothed_output, reference_output, rtol=rtol, atol=atol)