import torch
from torch._inductor.select_algorithm import extern_kernels
import triton
import triton.language as tl
from torch._inductor.runtime.triton_heuristics import grid
from torch._C import _cuda_getCurrentRawStream as get_raw_stream
from torch.nn import Module
import math
from torch.nn.parameter import Parameter
from torch.nn.modules.module import Module
from torch.nn.modules.utils import _pair
from torch.nn.functional import unfold
from torch.nn import Parameter
assert_size_stride = torch._C._dynamo.guards.assert_size_stride
empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda
reinterpret_tensor = torch._C._dynamo.guards._reinterpret_tensor


@triton.jit
def triton_poi_fused_im2col_0(in_ptr0, out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x3 = xindex
    tmp0 = tl.load(in_ptr0 + x3, xmask)
    tl.store(out_ptr0 + x3, tmp0, xmask)


@triton.jit
def triton_poi_fused_add_1(in_out_ptr0, in_ptr0, xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x0 = xindex % 4
    tmp0 = tl.load(in_out_ptr0 + x2, xmask)
    tmp1 = tl.load(in_ptr0 + x0, xmask, eviction_policy='evict_last')
    tmp2 = tmp0 + tmp1
    tl.store(in_out_ptr0 + x2, tmp2, xmask)


def call(args):
    primals_1, primals_2, primals_3 = args
    args.clear()
    assert_size_stride(primals_1, (1, 1, 4, 4, 4, 4), (256, 256, 64, 16, 4, 1))
    assert_size_stride(primals_2, (4, 1, 1), (1, 1, 1))
    assert_size_stride(primals_3, (4, 4, 4, 4), (64, 16, 4, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((4, 4, 4, 1, 4, 1), (64, 16, 4, 4, 1, 1),
            torch.float32)
        get_raw_stream(0)
        triton_poi_fused_im2col_0[grid(256)](primals_3, buf0, 256, XBLOCK=
            128, num_warps=4, num_stages=1)
        del primals_3
        buf1 = empty_strided_cuda((4, 1, 4), (4, 4, 1), torch.float32)
        extern_kernels.bmm(reinterpret_tensor(buf0, (4, 1, 64), (64, 0, 1),
            0), reinterpret_tensor(primals_1, (4, 64, 4), (0, 1, 64), 0),
            out=buf1)
        del primals_1
        buf2 = reinterpret_tensor(buf1, (4, 4, 1, 1), (4, 1, 4, 4), 0)
        del buf1
        triton_poi_fused_add_1[grid(16)](buf2, primals_2, 16, XBLOCK=16,
            num_warps=1, num_stages=1)
        del primals_2
    return buf2, reinterpret_tensor(buf0, (4, 64, 1), (64, 1, 4), 0)


def conv2d_local(input: 'torch.Tensor', weight: 'torch.Tensor', bias=None,
    padding: 'Pairable'=0, stride: 'Pairable'=1, dilation: 'Pairable'=1,
    data_format: 'str'='channels_first'):
    """Calculate the local convolution.

    Args:
        input:
        weight:
        bias:
        padding:
        stride:
        dilation:
        data_format: For Keras compatibility

    Returns:

    """
    if input.dim() != 4:
        raise NotImplementedError(
            'Input Error: Only 4D input Tensors supported (got {}D)'.format
            (input.dim()))
    if weight.dim() != 6:
        raise NotImplementedError(
            'Input Error: Only 6D weight Tensors supported (got {}D)'.
            format(weight.dim()))
    (out_height, out_width, out_channels, in_channels, kernel_height,
        kernel_width) = weight.size()
    kernel_size = kernel_height, kernel_width
    if data_format == 'channels_first':
        cols = unfold(input, kernel_size, dilation=dilation, padding=
            padding, stride=stride)
        reshaped_input = cols.view(cols.size(0), cols.size(1), cols.size(2), 1
            ).permute(0, 2, 3, 1)
    else:
        stride_y, stride_x = _pair(stride)
        feature_dim = in_channels * kernel_height * kernel_width
        xs = []
        for i in range(out_height):
            for j in range(out_width):
                y = i * stride_y
                slice_row = slice(y, y + kernel_size[0])
                x = j * stride_x
                slice_col = slice(x, x + kernel_size[1])
                val = input[:, slice_row, slice_col, :].contiguous()
                xs.append(val.view(input.shape[0], 1, -1, feature_dim))
        concated = torch.cat(xs, dim=1)
        reshaped_input = concated
    output_size = out_height * out_width
    input_size = in_channels * kernel_height * kernel_width
    weights_view = weight.view(output_size, out_channels, input_size)
    permuted_weights = weights_view.permute(0, 2, 1)
    out = torch.matmul(reshaped_input, permuted_weights)
    out = out.view(reshaped_input.shape[0], out_height, out_width, out_channels
        ).permute(0, 3, 1, 2)
    if data_format == 'channels_last':
        out = out.permute(0, 2, 3, 1)
    if bias is not None:
        final_bias = bias.expand_as(out)
        out = out + final_bias
    return out


class Conv2dLocalNew(Module):
    """A 2D locally connected layer.

    Attributes:
        weight (torch.Tensor): The weights. out_height x out_width x out_channels x in_channels x kernel_height x kernel_width
        kernel_size (Tuple[int, int]): The height and width of the convolutional kernels.
        stride (Tuple[int, int]): The stride height and width.
    """

    def __init__(self, in_height: 'int', in_width: 'int', in_channels:
        'int', out_channels: 'int', kernel_size: 'Pairable', stride:
        'Pairable'=1, padding: 'Pairable'=0, bias: 'bool'=True, dilation:
        'Pairable'=1, data_format='channels_first'):
        super(Conv2dLocalNew, self).__init__()
        self.data_format = data_format
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = _pair(kernel_size)
        self.stride = _pair(stride)
        self.padding = _pair(padding)
        self.dilation = _pair(dilation)
        self.in_height = in_height
        self.in_width = in_width
        self.out_height = int(math.floor((in_height + 2 * self.padding[0] -
            self.dilation[0] * (self.kernel_size[0] - 1) - 1) / self.stride
            [0] + 1))
        self.out_width = int(math.floor((in_width + 2 * self.padding[1] - 
            self.dilation[1] * (self.kernel_size[1] - 1) - 1) / self.stride
            [1] + 1))
        self.out_channels = out_channels
        self.weight = Parameter(torch.Tensor(self.out_height, self.
            out_width, out_channels, in_channels, *self.kernel_size))
        if bias:
            if self.data_format == 'channels_first':
                self.bias = Parameter(torch.Tensor(out_channels, self.
                    out_height, self.out_width))
            else:
                self.bias = Parameter(torch.Tensor(self.out_height, self.
                    out_width, out_channels))
        else:
            self.register_parameter('bias', None)
        self.reset_parameters()

    @property
    def input_shape(self):
        """The expected input shape for this module."""
        if self.data_format == 'channels_first':
            shape = self.in_channels, self.in_height, self.in_width
        else:
            shape = self.in_height, self.in_width, self.in_channels
        return torch.Tensor(shape)

    def reset_parameters(self):
        """Reset the parameters of the layer."""
        n = self.in_channels
        for k in self.kernel_size:
            n *= k
        stdv = 1.0 / math.sqrt(n)
        self.weight.data.uniform_(-stdv, stdv)
        if self.bias is not None:
            self.bias.data.uniform_(-stdv, stdv)

    def __repr__(self):
        s = (
            '{name}({in_channels}, {out_channels}, kernel_size={kernel_size}, stride={stride}'
            )
        if self.padding != (0,) * len(self.padding):
            s += ', padding={padding}'
        if self.dilation != (1,) * len(self.dilation):
            s += ', dilation={dilation}'
        if self.bias is None:
            s += ', bias=False'
        s += ')'
        return s.format(name=self.__class__.__name__, **self.__dict__)

    def forward(self, input_0):
        primals_1 = self.weight
        primals_2 = self.bias
        primals_3 = input_0
        output = call([primals_1, primals_2, primals_3])
        return output[0]


Measured latency_ms: 0.039558