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
import torch.nn.functional as F
import torch.utils.data
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):
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    tl.full([XBLOCK], True, tl.int1)
    x3 = xindex
    x1 = xindex // 262144 % 8
    tmp0 = tl.load(in_out_ptr0 + x3, None)
    tmp1 = tl.load(in_ptr0 + x1, None, eviction_policy='evict_last')
    tmp2 = tmp0 + tmp1
    tl.store(in_out_ptr0 + x3, tmp2, None)


def call(args):
    primals_1, primals_2, primals_3 = args
    args.clear()
    assert_size_stride(primals_1, (4, 1024, 64, 64, 64), (268435456, 262144,
        4096, 64, 1))
    assert_size_stride(primals_2, (8, 1024, 1, 1, 1), (1024, 1, 1, 1, 1))
    assert_size_stride(primals_3, (8,), (1,))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = extern_kernels.convolution(primals_1, primals_2, stride=(1, 
            1, 1), padding=(0, 0, 0), dilation=(1, 1, 1), transposed=False,
            output_padding=(0, 0, 0), groups=1, bias=None)
        assert_size_stride(buf0, (4, 8, 64, 64, 64), (2097152, 262144, 4096,
            64, 1))
        buf1 = buf0
        del buf0
        get_raw_stream(0)
        triton_poi_fused_convolution_0[grid(8388608)](buf1, primals_3, 
            8388608, XBLOCK=512, num_warps=8, num_stages=1)
        del primals_3
    return buf1, primals_1, primals_2


class Unit3D(nn.Module):
    """Basic unit containing Conv3D + BatchNorm + non-linearity."""

    def __init__(self, in_channels, output_channels, kernel_shape=(1, 1, 1),
        stride=(1, 1, 1), padding=0, activation_fn=F.relu, use_batch_norm=
        True, use_bias=False, name='unit_3d'):
        """Initializes Unit3D module."""
        super(Unit3D, self).__init__()
        self._output_channels = output_channels
        self._kernel_shape = kernel_shape
        self._stride = stride
        self._use_batch_norm = use_batch_norm
        self._activation_fn = activation_fn
        self._use_bias = use_bias
        self.name = name
        self.padding = padding
        self.conv3d = nn.Conv3d(in_channels=in_channels, out_channels=self.
            _output_channels, kernel_size=self._kernel_shape, stride=self.
            _stride, padding=0, bias=self._use_bias)
        if self._use_batch_norm:
            self.bn = nn.BatchNorm3d(self._output_channels, eps=0.001,
                momentum=0.01)

    def compute_pad(self, dim, s):
        """Get the zero padding number."""
        if s % self._stride[dim] == 0:
            return max(self._kernel_shape[dim] - self._stride[dim], 0)
        else:
            return max(self._kernel_shape[dim] - s % self._stride[dim], 0)

    def forward(self, x):
        """
        Connects the module to inputs. Dynamically pad based on input size in forward function.
        Args:
            x: Inputs to the Unit3D component.

        Returns:
            Outputs from the module.
        """
        _batch, _channel, time, height, width = x.size()
        pad_t = self.compute_pad(0, time)
        pad_h = self.compute_pad(1, height)
        pad_w = self.compute_pad(2, width)
        pad_t_front = pad_t // 2
        pad_t_back = pad_t - pad_t_front
        pad_h_front = pad_h // 2
        pad_h_back = pad_h - pad_h_front
        pad_w_front = pad_w // 2
        pad_w_back = pad_w - pad_w_front
        pad = (pad_w_front, pad_w_back, pad_h_front, pad_h_back,
            pad_t_front, pad_t_back)
        x = F.pad(x, pad)
        x = self.conv3d(x)
        if self._use_batch_norm:
            x = self.bn(x)
        if self._activation_fn is not None:
            x = self._activation_fn(x)
        return x


class ClassNetVideoConvNew(nn.Module):
    """Classifier network for video input refer to MMSADA.

    Args:
        input_size (int, optional): the dimension of the final feature vector. Defaults to 1024.
        n_class (int, optional): the number of classes. Defaults to 8.

    References:
        Munro Jonathan, and Dima Damen. "Multi-modal domain adaptation for fine-grained action recognition."
        In CVPR, pp. 122-132. 2020.
    """

    def __init__(self, input_size=1024, n_class=8):
        super(ClassNetVideoConvNew, self).__init__()
        self.dp = nn.Dropout()
        self.logits = Unit3D(in_channels=input_size, output_channels=
            n_class, kernel_shape=[1, 1, 1], padding=0, activation_fn=None,
            use_batch_norm=False, use_bias=True)

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


Measured latency_ms: 0.0242