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._inductor.runtime import triton_helpers
from torch._inductor.runtime.triton_helpers import math as tl_math
import math
import torch.utils.data
from torch import nn
from torch.nn.modules.utils import _pair
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_convolution_0(in_out_ptr0, in_ptr0, xnumel, XBLOCK: tl
    .constexpr):
    xnumel = 128
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x3 = xindex
    x1 = xindex // 16 % 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)


@triton.jit
def triton_per_fused__softmax_1(in_ptr0, out_ptr2, xnumel, rnumel, XBLOCK:
    tl.constexpr):
    xnumel = 64
    RBLOCK: tl.constexpr = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
    xmask = xindex < xnumel
    rindex = tl.arange(0, RBLOCK)[None, :]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    r1 = rindex
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + (r1 + 16 * x0), xmask, other=0.0)
    tmp1 = 1.0
    tmp2 = tmp0 * tmp1
    tmp3 = tl.broadcast_to(tmp2, [XBLOCK, RBLOCK])
    tmp5 = tl.where(xmask, tmp3, float('-inf'))
    tmp6 = triton_helpers.max2(tmp5, 1)[:, None]
    tmp7 = tmp2 - tmp6
    tmp8 = 1.414213562373095
    tmp9 = tmp7 * tmp8
    tmp10 = tl_math.exp(tmp9)
    tmp11 = tl.broadcast_to(tmp10, [XBLOCK, RBLOCK])
    tmp13 = tl.where(xmask, tmp11, 0)
    tmp14 = tl.sum(tmp13, 1)[:, None]
    tmp15 = tmp10 / tmp14
    tl.store(out_ptr2 + (r1 + 16 * x0), tmp15, xmask)


@triton.jit
def triton_poi_fused_add_convolution_2(in_ptr0, in_ptr1, in_ptr2, out_ptr0,
    ynumel, xnumel, YBLOCK: tl.constexpr, XBLOCK: tl.constexpr):
    ynumel = 16
    xnumel = 16
    yoffset = tl.program_id(1) * YBLOCK
    yindex = yoffset + tl.arange(0, YBLOCK)[None, :]
    ymask = yindex < ynumel
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
    xmask = xindex < xnumel
    x2 = xindex
    y3 = yindex
    y0 = yindex % 4
    y1 = yindex // 4
    tmp0 = tl.load(in_ptr0 + (x2 + 16 * y3), xmask & ymask)
    tmp1 = tl.load(in_ptr1 + (y0 + 4 * x2 + 64 * y1), xmask & ymask)
    tmp2 = tl.load(in_ptr2 + y0, ymask, eviction_policy='evict_last')
    tmp3 = tmp1 + tmp2
    tmp4 = tmp0 + tmp3
    tl.store(out_ptr0 + (x2 + 16 * y3), tmp4, xmask & ymask)


def call(args):
    (primals_1, primals_2, primals_3, primals_4, primals_5, primals_6,
        primals_7, primals_8, primals_9) = args
    args.clear()
    assert_size_stride(primals_1, (4, 4, 4, 4), (64, 16, 4, 1))
    assert_size_stride(primals_2, (2, 4, 1, 1), (4, 1, 1, 1))
    assert_size_stride(primals_3, (2,), (1,))
    assert_size_stride(primals_4, (2, 4, 1, 1), (4, 1, 1, 1))
    assert_size_stride(primals_5, (2,), (1,))
    assert_size_stride(primals_6, (2, 4, 1, 1), (4, 1, 1, 1))
    assert_size_stride(primals_7, (2,), (1,))
    assert_size_stride(primals_8, (4, 2, 1, 1), (2, 1, 1, 1))
    assert_size_stride(primals_9, (4,), (1,))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = extern_kernels.convolution(primals_1, primals_2, 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, 4, 4), (32, 16, 4, 1))
        buf1 = extern_kernels.convolution(primals_1, primals_4, stride=(1, 
            1), padding=(0, 0), dilation=(1, 1), transposed=False,
            output_padding=(0, 0), groups=1, bias=None)
        assert_size_stride(buf1, (4, 2, 4, 4), (32, 16, 4, 1))
        buf2 = extern_kernels.convolution(primals_1, primals_6, stride=(1, 
            1), padding=(0, 0), dilation=(1, 1), transposed=False,
            output_padding=(0, 0), groups=1, bias=None)
        assert_size_stride(buf2, (4, 2, 4, 4), (32, 16, 4, 1))
        buf3 = buf1
        del buf1
        get_raw_stream(0)
        triton_poi_fused_convolution_0[grid(128)](buf3, primals_5, 128,
            XBLOCK=128, num_warps=4, num_stages=1)
        del primals_5
        buf4 = buf2
        del buf2
        triton_poi_fused_convolution_0[grid(128)](buf4, primals_7, 128,
            XBLOCK=128, num_warps=4, num_stages=1)
        del primals_7
        buf5 = empty_strided_cuda((4, 16, 16), (256, 16, 1), torch.float32)
        extern_kernels.bmm(reinterpret_tensor(buf3, (4, 16, 2), (32, 1, 16),
            0), reinterpret_tensor(buf4, (4, 2, 16), (32, 16, 1), 0), out=buf5)
        buf8 = empty_strided_cuda((4, 16, 16), (256, 16, 1), torch.float32)
        triton_per_fused__softmax_1[grid(64)](buf5, buf8, 64, 16, XBLOCK=1,
            num_warps=2, num_stages=1)
        del buf5
        buf9 = buf0
        del buf0
        triton_poi_fused_convolution_0[grid(128)](buf9, primals_3, 128,
            XBLOCK=128, num_warps=4, num_stages=1)
        del primals_3
        buf10 = empty_strided_cuda((4, 16, 2), (32, 2, 1), torch.float32)
        extern_kernels.bmm(buf8, reinterpret_tensor(buf9, (4, 16, 2), (32, 
            1, 16), 0), out=buf10)
        buf11 = extern_kernels.convolution(reinterpret_tensor(buf10, (4, 2,
            4, 4), (32, 1, 8, 2), 0), primals_8, stride=(1, 1), padding=(0,
            0), dilation=(1, 1), transposed=False, output_padding=(0, 0),
            groups=1, bias=None)
        assert_size_stride(buf11, (4, 4, 4, 4), (64, 1, 16, 4))
        buf12 = empty_strided_cuda((4, 4, 4, 4), (64, 16, 4, 1), torch.float32)
        triton_poi_fused_add_convolution_2[grid(16, 16)](primals_1, buf11,
            primals_9, buf12, 16, 16, XBLOCK=16, YBLOCK=16, num_warps=4,
            num_stages=1)
        del buf11
        del primals_9
    return (buf12, primals_1, primals_2, primals_4, primals_6, primals_8,
        buf8, reinterpret_tensor(buf10, (4, 2, 4, 4), (32, 1, 8, 2), 0),
        reinterpret_tensor(buf9, (4, 2, 16), (32, 16, 1), 0),
        reinterpret_tensor(buf3, (4, 2, 16), (32, 16, 1), 0),
        reinterpret_tensor(buf4, (4, 16, 2), (32, 1, 16), 0))


def get_group_gn(dim, dim_per_gp, num_groups):
    """get number of groups used by GroupNorm, based on number of channels."""
    assert dim_per_gp == -1 or num_groups == -1, 'GroupNorm: can only specify G or C/G.'
    if dim_per_gp > 0:
        assert dim % dim_per_gp == 0, 'dim: {}, dim_per_gp: {}'.format(dim,
            dim_per_gp)
        group_gn = dim // dim_per_gp
    else:
        assert dim % num_groups == 0, 'dim: {}, num_groups: {}'.format(dim,
            num_groups)
        group_gn = num_groups
    return group_gn


def group_norm(out_channels, affine=True, divisor=1):
    out_channels = out_channels // divisor
    dim_per_gp = cfg.MODEL.GROUP_NORM.DIM_PER_GP // divisor
    num_groups = cfg.MODEL.GROUP_NORM.NUM_GROUPS // divisor
    eps = cfg.MODEL.GROUP_NORM.EPSILON
    return torch.nn.GroupNorm(get_group_gn(out_channels, dim_per_gp,
        num_groups), out_channels, eps, affine)


def conv_with_kaiming_uniform(use_gn=False, use_relu=False, use_deformable=
    False, use_bn=False):

    def make_conv(in_channels, out_channels, kernel_size, stride=1, dilation=1
        ):
        if use_deformable:
            conv_func = DFConv2d
        else:
            conv_func = Conv2d
        conv = conv_func(in_channels, out_channels, kernel_size=kernel_size,
            stride=stride, padding=dilation * (kernel_size - 1) // 2,
            dilation=dilation, bias=not (use_gn or use_bn))
        if not use_deformable:
            nn.init.kaiming_uniform_(conv.weight, a=1)
            if not (use_gn or use_bn):
                nn.init.constant_(conv.bias, 0)
        module = [conv]
        if use_gn:
            module.append(group_norm(out_channels))
        elif use_bn:
            module.append(nn.BatchNorm2d(out_channels))
        if use_relu:
            module.append(nn.ReLU(inplace=True))
        if len(module) > 1:
            return nn.Sequential(*module)
        return conv
    return make_conv


class DeformConv(nn.Module):

    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
        padding=0, dilation=1, groups=1, deformable_groups=1, bias=False):
        super(DeformConv, self).__init__()
        self.with_bias = bias
        assert in_channels % groups == 0, 'in_channels {} cannot be divisible by groups {}'.format(
            in_channels, groups)
        assert out_channels % groups == 0, 'out_channels {} cannot be divisible by groups {}'.format(
            out_channels, groups)
        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.groups = groups
        self.deformable_groups = deformable_groups
        self.weight = nn.Parameter(torch.Tensor(out_channels, in_channels //
            self.groups, *self.kernel_size))
        if self.with_bias:
            self.bias = nn.Parameter(torch.Tensor(out_channels))
        self.reset_parameters()

    def reset_parameters(self):
        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.with_bias:
            torch.nn.init.constant_(self.bias, 0.0)

    def forward(self, input, offset):
        y = deform_conv(input, offset, self.weight, self.stride, self.
            padding, self.dilation, self.groups, self.deformable_groups)
        if self.with_bias:
            assert len(y.size()) == 4
            y = y + self.bias.reshape(1, -1, 1, 1)
        return y

    def __repr__(self):
        return ''.join(['{}('.format(self.__class__.__name__),
            'in_channels={}, '.format(self.in_channels),
            'out_channels={}, '.format(self.out_channels),
            'kernel_size={}, '.format(self.kernel_size), 'stride={}, '.
            format(self.stride), 'dilation={}, '.format(self.dilation),
            'padding={}, '.format(self.padding), 'groups={}, '.format(self.
            groups), 'deformable_groups={}, '.format(self.deformable_groups
            ), 'bias={})'.format(self.with_bias)])


class ModulatedDeformConv(nn.Module):

    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
        padding=0, dilation=1, groups=1, deformable_groups=1, bias=True):
        super(ModulatedDeformConv, self).__init__()
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.kernel_size = _pair(kernel_size)
        self.stride = stride
        self.padding = _pair(padding)
        self.dilation = dilation
        self.groups = groups
        self.deformable_groups = deformable_groups
        self.with_bias = bias
        self.weight = nn.Parameter(torch.Tensor(out_channels, in_channels //
            groups, *self.kernel_size))
        if bias:
            self.bias = nn.Parameter(torch.Tensor(out_channels))
        else:
            self.register_parameter('bias', None)
        self.reset_parameters()

    def reset_parameters(self):
        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.zero_()

    def forward(self, input, offset, mask):
        return modulated_deform_conv(input, offset, mask, self.weight, self
            .bias, self.stride, self.padding, self.dilation, self.groups,
            self.deformable_groups)

    def __repr__(self):
        return ''.join(['{}('.format(self.__class__.__name__),
            'in_channels={}, '.format(self.in_channels),
            'out_channels={}, '.format(self.out_channels),
            'kernel_size={}, '.format(self.kernel_size), 'stride={}, '.
            format(self.stride), 'dilation={}, '.format(self.dilation),
            'padding={}, '.format(self.padding), 'groups={}, '.format(self.
            groups), 'deformable_groups={}, '.format(self.deformable_groups
            ), 'bias={})'.format(self.with_bias)])


class _NewEmptyTensorOp(torch.autograd.Function):

    @staticmethod
    def forward(ctx, x, new_shape):
        ctx.shape = x.shape
        return x.new_empty(new_shape)

    @staticmethod
    def backward(ctx, grad):
        shape = ctx.shape
        return _NewEmptyTensorOp.apply(grad, shape), None


class Conv2d(torch.nn.Conv2d):

    def forward(self, x):
        if x.numel() > 0:
            return super(Conv2d, self).forward(x)
        output_shape = [((i + 2 * p - (di * (k - 1) + 1)) // d + 1) for i,
            p, di, k, d in zip(x.shape[-2:], self.padding, self.dilation,
            self.kernel_size, self.stride)]
        output_shape = [x.shape[0], self.weight.shape[0]] + output_shape
        return _NewEmptyTensorOp.apply(x, output_shape)


class DFConv2d(nn.Module):
    """Deformable convolutional layer"""

    def __init__(self, in_channels, out_channels, with_modulated_dcn=True,
        kernel_size=3, stride=1, groups=1, dilation=1, deformable_groups=1,
        bias=False, padding=None):
        super(DFConv2d, self).__init__()
        if isinstance(kernel_size, (list, tuple)):
            assert isinstance(stride, (list, tuple))
            assert isinstance(dilation, (list, tuple))
            assert len(kernel_size) == 2
            assert len(stride) == 2
            assert len(dilation) == 2
            padding = dilation[0] * (kernel_size[0] - 1) // 2, dilation[1] * (
                kernel_size[1] - 1) // 2
            offset_base_channels = kernel_size[0] * kernel_size[1]
        else:
            padding = dilation * (kernel_size - 1) // 2
            offset_base_channels = kernel_size * kernel_size
        if with_modulated_dcn:
            offset_channels = offset_base_channels * 3
            conv_block = ModulatedDeformConv
        else:
            offset_channels = offset_base_channels * 2
            conv_block = DeformConv
        self.offset = Conv2d(in_channels, deformable_groups *
            offset_channels, kernel_size=kernel_size, stride=stride,
            padding=padding, groups=1, dilation=dilation)
        for l in [self.offset]:
            nn.init.kaiming_uniform_(l.weight, a=1)
            torch.nn.init.constant_(l.bias, 0.0)
        self.conv = conv_block(in_channels, out_channels, kernel_size=
            kernel_size, stride=stride, padding=padding, dilation=dilation,
            groups=groups, deformable_groups=deformable_groups, bias=bias)
        self.with_modulated_dcn = with_modulated_dcn
        self.kernel_size = kernel_size
        self.stride = stride
        self.padding = padding
        self.dilation = dilation
        self.offset_split = offset_base_channels * deformable_groups * 2

    def forward(self, x, return_offset=False):
        if x.numel() > 0:
            if not self.with_modulated_dcn:
                offset_mask = self.offset(x)
                x = self.conv(x, offset_mask)
            else:
                offset_mask = self.offset(x)
                offset = offset_mask[:, :self.offset_split, :, :]
                mask = offset_mask[:, self.offset_split:, :, :].sigmoid()
                x = self.conv(x, offset, mask)
            if return_offset:
                return x, offset_mask
            return x
        output_shape = [((i + 2 * p - (di * (k - 1) + 1)) // d + 1) for i,
            p, di, k, d in zip(x.shape[-2:], self.padding, self.dilation,
            self.kernel_size, self.stride)]
        output_shape = [x.shape[0], self.conv.weight.shape[0]] + output_shape
        return _NewEmptyTensorOp.apply(x, output_shape)


class NonLocal2DNew(nn.Module):
    """Non-local module.
    See https://arxiv.org/abs/1711.07971 for details.
    Args:
        in_channels (int): Channels of the input feature map.
        reduction (int): Channel reduction ratio.
        use_scale (bool): Whether to scale pairwise_weight by 1/inter_channels.
        mode (str): Options are `embedded_gaussian` and `dot_product`.
    """

    def __init__(self, in_channels, reduction=2, use_scale=True, use_gn=
        False, use_deformable=False, mode='embedded_gaussian'):
        super(NonLocal2DNew, self).__init__()
        self.in_channels = in_channels
        self.reduction = reduction
        self.use_scale = use_scale
        self.inter_channels = in_channels // reduction
        self.mode = mode
        assert mode in ['embedded_gaussian', 'dot_product']
        ConvModule = conv_with_kaiming_uniform()
        last_conv = conv_with_kaiming_uniform(use_gn=use_gn, use_deformable
            =use_deformable)
        self.g = ConvModule(self.in_channels, self.inter_channels,
            kernel_size=1)
        self.theta = ConvModule(self.in_channels, self.inter_channels,
            kernel_size=1)
        self.phi = ConvModule(self.in_channels, self.inter_channels,
            kernel_size=1)
        self.conv_out = last_conv(self.inter_channels, self.in_channels,
            kernel_size=1)

    def embedded_gaussian(self, theta_x, phi_x):
        pairwise_weight = torch.matmul(theta_x, phi_x)
        if self.use_scale:
            pairwise_weight /= theta_x.shape[-1] ** -0.5
        pairwise_weight = pairwise_weight.softmax(dim=-1)
        return pairwise_weight

    def dot_product(self, theta_x, phi_x):
        pairwise_weight = torch.matmul(theta_x, phi_x)
        pairwise_weight /= pairwise_weight.shape[-1]
        return pairwise_weight

    def forward(self, input_0):
        primals_2 = self.g.weight
        primals_3 = self.g.bias
        primals_4 = self.theta.weight
        primals_5 = self.theta.bias
        primals_6 = self.phi.weight
        primals_7 = self.phi.bias
        primals_8 = self.conv_out.weight
        primals_9 = self.conv_out.bias
        primals_1 = input_0
        output = call([primals_1, primals_2, primals_3, primals_4,
            primals_5, primals_6, primals_7, primals_8, primals_9])
        return output[0]


Measured latency_ms: 0.121871