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 import nn
from torchvision import models as models
import torch.nn.parallel
import torch.optim
import torch.utils.data
import torch.utils.data.distributed
from torchvision.transforms import *
import torch.onnx
assert_size_stride = torch._C._dynamo.guards.assert_size_stride


@triton.jit
def triton_poi_fused_convolution_0(in_out_ptr0, in_ptr0, xnumel, XBLOCK: tl
    .constexpr):
    xnumel = 31752
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x3 = xindex
    x1 = xindex // 3969 % 2
    tmp0 = tl.load(in_out_ptr0 + x3, xmask)
    tmp1 = tl.load(in_ptr0 + x1, xmask, eviction_policy='evict_last')
    tmp2 = tmp0 + tmp1
    tl.store(in_out_ptr0 + x3, tmp2, xmask)


def call(args):
    primals_1, primals_2, primals_3 = args
    args.clear()
    assert_size_stride(primals_1, (2, 1, 2, 2), (4, 4, 2, 1))
    assert_size_stride(primals_2, (2,), (1,))
    assert_size_stride(primals_3, (4, 1, 64, 64), (4096, 4096, 64, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = extern_kernels.convolution(primals_3, primals_1, stride=(1, 
            1), padding=(0, 0), dilation=(1, 1), transposed=False,
            output_padding=(0, 0), groups=1, bias=None)
        assert_size_stride(buf0, (4, 2, 63, 63), (7938, 3969, 63, 1))
        buf1 = buf0
        del buf0
        get_raw_stream(0)
        triton_poi_fused_convolution_0[grid(31752)](buf1, primals_2, 31752,
            XBLOCK=256, num_warps=4, num_stages=1)
        del primals_2
    return buf1, primals_1, primals_3


def fill_bias(module, value):
    module.bias.data.fill_(value)


def fill_conv_weight(conv, value):
    conv.weight.data.fill_(value)
    with torch.no_grad():
        mask = torch.eye(conv.kernel_size[0])
        conv.weight += mask


def create_conv(in_channels, out_channels, kernel_size, weight_init, bias_init
    ):
    conv = nn.Conv2d(in_channels, out_channels, kernel_size)
    fill_conv_weight(conv, weight_init)
    fill_bias(conv, bias_init)
    return conv


class conv(nn.Module):

    def __init__(self, in_channels, out_channels, stride=1, groups=1):
        super(conv, self).__init__()
        self.conv1 = nn.Conv3d(in_channels=in_channels, out_channels=
            out_channels, kernel_size=(5, 5, 3), stride=(stride, stride, 1),
            padding=(2, 2, 1), bias=False, groups=groups)

    def forward(self, x):
        out = x
        out = self.conv1(out)
        return out


class BasicConvTestModelNew(nn.Module):

    def __init__(self, in_channels=1, out_channels=2, kernel_size=2,
        weight_init=-1, bias_init=-2):
        super().__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = kernel_size
        self.weight_init = weight_init
        self.bias_init = bias_init
        self.conv = create_conv(in_channels, out_channels, kernel_size,
            weight_init, bias_init)

    @staticmethod
    def default_weight():
        return torch.tensor([[[[0.0, -1.0], [-1.0, 0.0]]], [[[0.0, -1.0], [
            -1.0, 0.0]]]])

    @staticmethod
    def default_bias():
        return torch.tensor([-2.0, -2])

    @property
    def weights_num(self):
        return self.out_channels * self.kernel_size ** 2

    @property
    def bias_num(self):
        return self.kernel_size

    @property
    def nz_weights_num(self):
        return self.kernel_size * self.out_channels

    @property
    def nz_bias_num(self):
        return self.kernel_size

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


Measured latency_ms: 0.06345