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
import torch.nn as nn
assert_size_stride = torch._C._dynamo.guards.assert_size_stride
empty_strided_cuda = torch._C._dynamo.guards._empty_strided_cuda


@triton.jit
def triton_poi_fused_replication_pad2d_0(in_ptr0, out_ptr0, xnumel, XBLOCK:
    tl.constexpr):
    xnumel = 144
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 6
    x1 = xindex // 6 % 6
    x2 = xindex // 36
    x3 = xindex
    tmp0 = tl.load(in_ptr0 + (4 * (3 * (3 <= 0 * (0 >= -1 + x1) + (-1 + x1) *
        (-1 + x1 > 0)) + (0 * (0 >= -1 + x1) + (-1 + x1) * (-1 + x1 > 0)) *
        (0 * (0 >= -1 + x1) + (-1 + x1) * (-1 + x1 > 0) < 3)) + 16 * x2 + (
        3 * (3 <= 0 * (0 >= -1 + x0) + (-1 + x0) * (-1 + x0 > 0)) + (0 * (0 >=
        -1 + x0) + (-1 + x0) * (-1 + x0 > 0)) * (0 * (0 >= -1 + x0) + (-1 +
        x0) * (-1 + x0 > 0) < 3))), xmask, eviction_policy='evict_last')
    tl.store(out_ptr0 + x3, tmp0, xmask)


@triton.jit
def triton_poi_fused_convolution_leaky_relu_1(in_ptr0, in_ptr1, out_ptr0,
    out_ptr1, xnumel, XBLOCK: tl.constexpr):
    xnumel = 64
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + x0, xmask)
    tmp1 = tl.load(in_ptr1 + 0)
    tmp2 = tl.broadcast_to(tmp1, [XBLOCK])
    tmp3 = tmp0 + tmp2
    tmp4 = 0.0
    tmp5 = tmp3 > tmp4
    tmp6 = 0.01
    tmp7 = tmp3 * tmp6
    tmp8 = tl.where(tmp5, tmp3, tmp7)
    tl.store(out_ptr0 + x0, tmp5, xmask)
    tl.store(out_ptr1 + x0, tmp8, xmask)


@triton.jit
def triton_poi_fused_max_pool2d_with_indices_2(in_ptr0, out_ptr0, out_ptr1,
    xnumel, XBLOCK: tl.constexpr):
    xnumel = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 2
    x1 = xindex // 2
    x2 = xindex
    tmp0 = tl.load(in_ptr0 + (2 * x0 + 8 * x1), xmask, eviction_policy=
        'evict_last')
    tmp1 = tl.load(in_ptr0 + (1 + 2 * x0 + 8 * x1), xmask, eviction_policy=
        'evict_last')
    tmp3 = tl.load(in_ptr0 + (4 + 2 * x0 + 8 * x1), xmask, eviction_policy=
        'evict_last')
    tmp5 = tl.load(in_ptr0 + (5 + 2 * x0 + 8 * x1), xmask, eviction_policy=
        'evict_last')
    tmp2 = triton_helpers.maximum(tmp1, tmp0)
    tmp4 = triton_helpers.maximum(tmp3, tmp2)
    tmp6 = triton_helpers.maximum(tmp5, tmp4)
    tmp7 = tmp1 > tmp0
    tmp8 = tl.full([1], 1, tl.int8)
    tmp9 = tl.full([1], 0, tl.int8)
    tmp10 = tl.where(tmp7, tmp8, tmp9)
    tmp11 = tmp3 > tmp2
    tmp12 = tl.full([1], 2, tl.int8)
    tmp13 = tl.where(tmp11, tmp12, tmp10)
    tmp14 = tmp5 > tmp4
    tmp15 = tl.full([1], 3, tl.int8)
    tmp16 = tl.where(tmp14, tmp15, tmp13)
    tl.store(out_ptr0 + x2, tmp6, xmask)
    tl.store(out_ptr1 + x2, tmp16, xmask)


def call(args):
    primals_1, primals_2, primals_3, primals_4, primals_5 = args
    args.clear()
    assert_size_stride(primals_1, (1, 1, 3, 3), (9, 9, 3, 1))
    assert_size_stride(primals_2, (1,), (1,))
    assert_size_stride(primals_3, (4, 1, 4, 4), (16, 16, 4, 1))
    assert_size_stride(primals_4, (1, 1, 3, 3), (9, 9, 3, 1))
    assert_size_stride(primals_5, (1,), (1,))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((4, 1, 6, 6), (36, 36, 6, 1), torch.float32)
        get_raw_stream(0)
        triton_poi_fused_replication_pad2d_0[grid(144)](primals_3, buf0, 
            144, XBLOCK=128, num_warps=4, num_stages=1)
        del primals_3
        buf1 = extern_kernels.convolution(buf0, primals_1, stride=(1, 1),
            padding=(0, 0), dilation=(1, 1), transposed=False,
            output_padding=(0, 0), groups=1, bias=None)
        assert_size_stride(buf1, (4, 1, 4, 4), (16, 16, 4, 1))
        buf2 = empty_strided_cuda((4, 1, 4, 4), (16, 16, 4, 1), torch.bool)
        buf3 = empty_strided_cuda((4, 1, 4, 4), (16, 16, 4, 1), torch.float32)
        triton_poi_fused_convolution_leaky_relu_1[grid(64)](buf1, primals_2,
            buf2, buf3, 64, XBLOCK=64, num_warps=1, num_stages=1)
        del primals_2
        buf4 = empty_strided_cuda((4, 1, 6, 6), (36, 36, 6, 1), torch.float32)
        triton_poi_fused_replication_pad2d_0[grid(144)](buf3, buf4, 144,
            XBLOCK=128, num_warps=4, num_stages=1)
        buf5 = extern_kernels.convolution(buf4, primals_4, stride=(1, 1),
            padding=(0, 0), dilation=(1, 1), transposed=False,
            output_padding=(0, 0), groups=1, bias=None)
        assert_size_stride(buf5, (4, 1, 4, 4), (16, 16, 4, 1))
        buf6 = empty_strided_cuda((4, 1, 4, 4), (16, 16, 4, 1), torch.bool)
        buf7 = buf1
        del buf1
        triton_poi_fused_convolution_leaky_relu_1[grid(64)](buf5, primals_5,
            buf6, buf7, 64, XBLOCK=64, num_warps=1, num_stages=1)
        del buf5
        del primals_5
        buf8 = empty_strided_cuda((4, 1, 2, 2), (4, 4, 2, 1), torch.float32)
        buf9 = empty_strided_cuda((4, 1, 2, 2), (4, 4, 2, 1), torch.int8)
        triton_poi_fused_max_pool2d_with_indices_2[grid(16)](buf7, buf8,
            buf9, 16, XBLOCK=16, num_warps=1, num_stages=1)
    return buf8, primals_1, primals_4, buf0, buf2, buf3, buf4, buf6, buf7, buf9


class ConvBlock(nn.Module):
    """ Conv layer block """

    def __init__(self, kernel, in_depth, conv_depth, stride=1, padding=1,
        normalization=False, norm_type='BN', pooling=False,
        bias_initialization='zeros', activation=True, dilation=1,
        return_before_pooling=False):
        """ ConvBlock constructor

    Args:
      kernel: kernel size (int)
      in_depth: depth of input tensor
      conv_depth: number of out channels produced by the convolution
      stride: stide of the convolution; default is 1.
      padding: zero-padding added to both sides before the convolution
        operation; default is 1.
      normalization: boolean flag to apply normalization after the conv;
        default is false.
      norm_type: normalization operation: 'BN' for batch-norm (default),
        'IN' for instance normalization.
      pooling: boolean flag to apply a 2 x 2 max-pooling with stride of 2
        before returning the final result; default is false.
      bias_initialization: bias initialization: 'zeros' (default) or 'ones'.
      activation: boolean flag to apply a leaky ReLU activation; default is
        true.
      dilation: spacing between conv kernel elements; default is 1.
      return_before_pooling: boolean flag to return the tensor before
        applying max-pooling (if 'pooling' is true); default is false.

    Returns:
      ConvBlock object with the selected settings.
    """
        super(ConvBlock, self).__init__()
        conv = torch.nn.Conv2d(in_depth, conv_depth, kernel, stride=stride,
            dilation=dilation, padding=padding, padding_mode='replicate')
        torch.nn.init.kaiming_normal_(conv.weight)
        if bias_initialization == 'ones':
            torch.nn.init.ones_(conv.bias)
        elif bias_initialization == 'zeros':
            torch.nn.init.zeros_(conv.bias)
        else:
            raise NotImplementedError
        if activation:
            self.activation = torch.nn.LeakyReLU(inplace=False)
        else:
            self.activation = None
        if normalization:
            if norm_type == 'BN':
                self.normalization = torch.nn.BatchNorm2d(conv_depth,
                    affine=True)
            elif norm_type == 'IN':
                self.normalization = torch.nn.InstanceNorm2d(conv_depth,
                    affine=False)
            else:
                raise NotImplementedError
        else:
            self.normalization = None
        self.conv = conv
        if pooling:
            self.pooling = torch.nn.MaxPool2d(2, stride=2)
        else:
            self.pooling = None
        self.return_before_pooling = return_before_pooling

    def forward(self, x):
        """ Forward function of ConvBlock module

    Args:
      x: input tensor.

    Returns:
      y: processed tensor.
    """
        x = self.conv(x)
        if self.normalization is not None:
            x = self.normalization(x)
        if self.activation is not None:
            x = self.activation(x)
        if self.pooling is not None:
            y = self.pooling(x)
        else:
            y = x
        if self.return_before_pooling:
            return y, x
        else:
            return y


class DoubleConvBlockNew(nn.Module):
    """ Double conv layers block """

    def __init__(self, in_depth, out_depth, mid_depth=None, kernel=3,
        stride=1, padding=None, dilation=None, normalization=False,
        norm_type='BN', pooling=True, return_before_pooling=False,
        normalization_block='Both'):
        """ DoubleConvBlock constructor

    Args:
      in_depth: depth of input tensor
      out_depth: number of out channels produced by the second convolution
      mid_depth: number of out channels produced by the first convolution;
        default is mid_depth = out_depth.
      kernel: kernel size (int); default is 3.
      stride: stide of the convolution; default is 1.
      padding: zero-padding added to both sides before the convolution
        operations; default is [1, 1].
      dilation: spacing between elements of each conv kernel; default is [1, 1].
      normalization: boolean flag to apply normalization after the conv;
        default is false.
      norm_type: normalization operation: 'BN' for batch-norm (default),
        'IN' for instance normalization.
      pooling: boolean flag to apply a 2 x 2 max-pooling with stride of 2
        before returning the final result; default is false.
      return_before_pooling: boolean flag to return the tensor before
        applying max-pooling (if 'pooling' is true); default is false.
      normalization_block: if normalization flag is set to true; this
        variable controls when to apply the normalization process. It can be:
        'Both' (apply normalization after both conv layers), 'First', or
        'Second'.

    Returns:
      DoubleConvBlock object with the selected settings.
    """
        super().__init__()
        if padding is None:
            padding = [1, 1]
        if dilation is None:
            dilation = [1, 1]
        if mid_depth is None:
            mid_depth = out_depth
        if normalization:
            if normalization_block == 'First':
                norm = [True, False]
            elif normalization_block == 'Second':
                norm = [False, True]
            elif normalization_block == 'Both':
                norm = [True, True]
            else:
                raise NotImplementedError
        else:
            norm = [False, False]
        self.double_conv_1 = ConvBlock(kernel=kernel, in_depth=in_depth,
            conv_depth=mid_depth, stride=stride, padding=padding[0],
            pooling=False, dilation=dilation[0], norm_type=norm_type,
            normalization=norm[0])
        self.double_conv_2 = ConvBlock(kernel=kernel, in_depth=mid_depth,
            conv_depth=out_depth, stride=stride, padding=padding[1],
            pooling=pooling, dilation=dilation[1], norm_type=norm_type,
            normalization=norm[1], return_before_pooling=return_before_pooling)

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


Measured latency_ms: 0.013863