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
import torch.nn
import torch.nn.functional
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_gather_repeat_1(in_ptr0, in_ptr1, out_ptr0, xnumel,
    XBLOCK: tl.constexpr):
    xnumel = 192
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x0 = xindex % 4
    x1 = xindex // 4 % 4
    x2 = xindex // 16 % 3
    x3 = xindex // 48
    x4 = xindex % 16
    x5 = xindex
    tmp0 = tl.load(in_ptr0 + (2 * (x1 % 2) + 4 * x2 + x0 % 2), xmask,
        eviction_policy='evict_last')
    tmp1 = tl.full([XBLOCK], 5, tl.int32)
    tmp2 = tmp0 + tmp1
    tmp3 = tmp0 < 0
    tmp4 = tl.where(tmp3, tmp2, tmp0)
    tl.device_assert((0 <= tmp4) & (tmp4 < 5) | ~xmask,
        'index out of bounds: 0 <= tmp4 < 5')
    tmp6 = tl.load(in_ptr1 + (x4 + 16 * tmp4 + 80 * x3), xmask)
    tl.store(out_ptr0 + x5, tmp6, xmask)


def call(args):
    arg0_1, arg1_1, arg2_1 = args
    args.clear()
    assert_size_stride(arg0_1, (4, 1, 4, 4), (16, 16, 4, 1))
    assert_size_stride(arg1_1, (5, 1, 3, 3), (9, 9, 3, 1))
    assert_size_stride(arg2_1, (1, 3, 2, 2), (12, 4, 2, 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)](arg0_1, buf0, 144,
            XBLOCK=256, num_warps=4, num_stages=1)
        del arg0_1
        buf1 = extern_kernels.convolution(buf0, arg1_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, 5, 4, 4), (80, 16, 4, 1))
        del arg1_1
        del buf0
        buf2 = empty_strided_cuda((4, 3, 4, 4), (48, 16, 4, 1), torch.float32)
        triton_poi_fused_gather_repeat_1[grid(192)](arg2_1, buf1, buf2, 192,
            XBLOCK=128, num_warps=4, num_stages=1)
        del arg2_1
        del buf1
    return buf2,


class Debayer3x3New(torch.nn.Module):
    """Demosaicing of Bayer images using 3x3 convolutions.

    Requires BG-Bayer color filter array layout. That is,
    the image[1,1]='B', image[1,2]='G'. This corresponds
    to OpenCV naming conventions.

    Compared to Debayer2x2 this method does not use upsampling.
    Instead, we identify five 3x3 interpolation kernels that
    are sufficient to reconstruct every color channel at every
    pixel location.

    We convolve the image with these 5 kernels using stride=1
    and a one pixel replication padding. Finally, we gather
    the correct channel values for each pixel location. Todo so,
    we recognize that the Bayer pattern repeats horizontally and
    vertically every 2 pixels. Therefore, we define the correct
    index lookups for a 2x2 grid cell and then repeat to image
    dimensions.

    Note, in every 2x2 grid cell we have red, blue and two greens 
    (G1,G2). The lookups for the two greens differ.
    """

    def __init__(self):
        super(Debayer3x3New, self).__init__()
        self.kernels = torch.nn.Parameter(torch.tensor([[0, 0, 0], [0, 1, 0
            ], [0, 0, 0], [0, 0.25, 0], [0.25, 0, 0.25], [0, 0.25, 0], [
            0.25, 0, 0.25], [0, 0, 0], [0.25, 0, 0.25], [0, 0, 0], [0.5, 0,
            0.5], [0, 0, 0], [0, 0.5, 0], [0, 0, 0], [0, 0.5, 0]]).view(5, 
            1, 3, 3), requires_grad=False)
        self.index = torch.nn.Parameter(torch.tensor([[0, 3], [4, 2], [1, 0
            ], [0, 1], [2, 4], [3, 0]]).view(1, 3, 2, 2), requires_grad=False)

    def forward(self, input_0):
        arg1_1 = self.kernels
        arg2_1 = self.index
        arg0_1 = input_0
        output = call([arg0_1, arg1_1, arg2_1])
        return output[0]


Measured latency_ms: 0.036443