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 as nn
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_replication_pad3d_0(in_ptr0, out_ptr0, xnumel, XBLOCK:
    tl.constexpr):
    xnumel = 384
    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 % 6
    x3 = xindex // 96
    x4 = xindex
    tmp0 = tl.load(in_ptr0 + (4 * (3 * (3 <= x1) + x1 * (x1 < 3)) + 16 * (3 *
        (3 <= 0 * (0 >= -1 + x2) + (-1 + x2) * (-1 + x2 > 0)) + (0 * (0 >= 
        -1 + x2) + (-1 + x2) * (-1 + x2 > 0)) * (0 * (0 >= -1 + x2) + (-1 +
        x2) * (-1 + x2 > 0) < 3)) + 64 * x3 + (3 * (3 <= x0) + x0 * (x0 < 3
        ))), xmask)
    tl.store(out_ptr0 + x4, tmp0, xmask)


@triton.jit
def triton_poi_fused_convolution_1(in_out_ptr0, in_ptr0, xnumel, XBLOCK: tl
    .constexpr):
    xnumel = 108
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x2 = xindex
    x1 = xindex // 27
    tmp0 = tl.load(in_out_ptr0 + x2, xmask)
    tmp1 = tl.load(in_ptr0 + x1, 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, (4, 4, 4, 4), (64, 16, 4, 1))
    assert_size_stride(primals_2, (4, 4, 4, 4, 4), (256, 64, 16, 4, 1))
    assert_size_stride(primals_3, (4,), (1,))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((4, 6, 4, 4), (96, 16, 4, 1), torch.float32)
        get_raw_stream(0)
        triton_poi_fused_replication_pad3d_0[grid(384)](primals_1, buf0, 
            384, XBLOCK=128, num_warps=4, num_stages=1)
        del primals_1
        buf1 = extern_kernels.convolution(reinterpret_tensor(buf0, (1, 4, 6,
            4, 4), (0, 96, 16, 4, 1), 0), primals_2, stride=(1, 1, 1),
            padding=(0, 1, 1), dilation=(1, 1, 1), transposed=False,
            output_padding=(0, 0, 0), groups=1, bias=None)
        assert_size_stride(buf1, (1, 4, 3, 3, 3), (108, 27, 9, 3, 1))
        buf2 = buf1
        del buf1
        triton_poi_fused_convolution_1[grid(108)](buf2, primals_3, 108,
            XBLOCK=128, num_warps=4, num_stages=1)
        del primals_3
    return reinterpret_tensor(buf2, (4, 3, 3, 3), (27, 9, 3, 1), 0
        ), primals_2, reinterpret_tensor(buf0, (1, 4, 6, 4, 4), (384, 96, 
        16, 4, 1), 0)


class MyConv3dNew(nn.Module):

    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
        bias=True):
        super(MyConv3dNew, self).__init__()
        self.kernel_size = kernel_size
        self.conv = nn.Conv3d(in_channels=in_channels, out_channels=
            out_channels, kernel_size=kernel_size, padding=(0, int((
            kernel_size - 1) / 2), int((kernel_size - 1) / 2)), bias=bias)

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


Measured latency_ms: 0.036365