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 logging
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__to_copy_clamp_mul_0(in_ptr0, out_ptr0, xnumel, XBLOCK:
    tl.constexpr):
    xnumel = 256
    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 = 128.0
    tmp2 = tmp0 * tmp1
    tmp3 = 0.0
    tmp4 = triton_helpers.maximum(tmp2, tmp3)
    tmp5 = 127.0
    tmp6 = triton_helpers.minimum(tmp4, tmp5)
    tmp7 = tmp6.to(tl.int64)
    tl.store(out_ptr0 + x0, tmp7, xmask)


@triton.jit
def triton_poi_fused_add_embedding_1(in_ptr0, in_ptr1, in_ptr2, out_ptr0,
    xnumel, XBLOCK: tl.constexpr):
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    tl.full([XBLOCK], True, tl.int1)
    x1 = xindex // 128 % 4
    x2 = xindex // 512
    x0 = xindex % 128
    x4 = xindex
    tmp0 = tl.load(in_ptr0 + (x1 + 16 * x2), None, eviction_policy='evict_last'
        )
    tmp7 = tl.load(in_ptr0 + (4 + x1 + 16 * x2), None, eviction_policy=
        'evict_last')
    tmp14 = tl.load(in_ptr0 + (8 + x1 + 16 * x2), None, eviction_policy=
        'evict_last')
    tmp21 = tl.load(in_ptr0 + (12 + x1 + 16 * x2), None, eviction_policy=
        'evict_last')
    tmp1 = tl.full([XBLOCK], 128, tl.int32)
    tmp2 = tmp0 + tmp1
    tmp3 = tmp0 < 0
    tmp4 = tl.where(tmp3, tmp2, tmp0)
    tl.device_assert((0 <= tmp4) & (tmp4 < 128),
        'index out of bounds: 0 <= tmp4 < 128')
    tmp6 = tl.load(in_ptr1 + (x0 + 128 * tmp4), None)
    tmp8 = tmp7 + tmp1
    tmp9 = tmp7 < 0
    tmp10 = tl.where(tmp9, tmp8, tmp7)
    tl.device_assert((0 <= tmp10) & (tmp10 < 128),
        'index out of bounds: 0 <= tmp10 < 128')
    tmp12 = tl.load(in_ptr2 + (x0 + 128 * tmp10), None)
    tmp13 = tmp6 + tmp12
    tmp15 = tmp14 + tmp1
    tmp16 = tmp14 < 0
    tmp17 = tl.where(tmp16, tmp15, tmp14)
    tl.device_assert((0 <= tmp17) & (tmp17 < 128),
        'index out of bounds: 0 <= tmp17 < 128')
    tmp19 = tl.load(in_ptr1 + (x0 + 128 * tmp17), None)
    tmp20 = tmp13 + tmp19
    tmp22 = tmp21 + tmp1
    tmp23 = tmp21 < 0
    tmp24 = tl.where(tmp23, tmp22, tmp21)
    tl.device_assert((0 <= tmp24) & (tmp24 < 128),
        'index out of bounds: 0 <= tmp24 < 128')
    tmp26 = tl.load(in_ptr2 + (x0 + 128 * tmp24), None)
    tmp27 = tmp20 + tmp26
    tl.store(out_ptr0 + x4, tmp27, None)


@triton.jit
def triton_poi_fused_relu_threshold_backward_2(in_out_ptr0, in_ptr0,
    out_ptr0, xnumel, XBLOCK: tl.constexpr):
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    tl.full([XBLOCK], True, tl.int1)
    x2 = xindex
    x0 = xindex % 128
    tmp0 = tl.load(in_out_ptr0 + x2, None)
    tmp1 = tl.load(in_ptr0 + x0, None, eviction_policy='evict_last')
    tmp2 = tmp0 + tmp1
    tmp3 = tl.full([1], 0, tl.int32)
    tmp4 = triton_helpers.maximum(tmp3, tmp2)
    tmp5 = 0.0
    tmp6 = tmp4 <= tmp5
    tl.store(in_out_ptr0 + x2, tmp4, None)
    tl.store(out_ptr0 + x2, tmp6, None)


def call(args):
    primals_1, primals_2, primals_3, primals_4, primals_5 = args
    args.clear()
    assert_size_stride(primals_1, (4, 4, 4, 4), (64, 16, 4, 1))
    assert_size_stride(primals_2, (128, 128), (128, 1))
    assert_size_stride(primals_3, (128, 128), (128, 1))
    assert_size_stride(primals_4, (128, 128), (128, 1))
    assert_size_stride(primals_5, (128,), (1,))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((4, 4, 4, 4), (64, 16, 4, 1), torch.int64)
        get_raw_stream(0)
        triton_poi_fused__to_copy_clamp_mul_0[grid(256)](primals_1, buf0, 
            256, XBLOCK=128, num_warps=4, num_stages=1)
        del primals_1
        buf1 = empty_strided_cuda((4, 4, 4, 128), (2048, 512, 128, 1),
            torch.float32)
        triton_poi_fused_add_embedding_1[grid(8192)](buf0, primals_2,
            primals_3, buf1, 8192, XBLOCK=256, num_warps=4, num_stages=1)
        del primals_2
        del primals_3
        buf2 = empty_strided_cuda((64, 128), (128, 1), torch.float32)
        extern_kernels.mm(reinterpret_tensor(buf1, (64, 128), (128, 1), 0),
            reinterpret_tensor(primals_4, (128, 128), (1, 128), 0), out=buf2)
        buf3 = reinterpret_tensor(buf2, (4, 4, 4, 128), (2048, 512, 128, 1), 0)
        del buf2
        buf4 = empty_strided_cuda((4, 4, 4, 128), (2048, 512, 128, 1),
            torch.bool)
        triton_poi_fused_relu_threshold_backward_2[grid(8192)](buf3,
            primals_5, buf4, 8192, XBLOCK=256, num_warps=4, num_stages=1)
        del primals_5
    return buf3, reinterpret_tensor(buf0, (4, 4, 4), (64, 16, 1), 0
        ), reinterpret_tensor(buf0, (4, 4, 4), (64, 16, 1), 4
        ), reinterpret_tensor(buf0, (4, 4, 4), (64, 16, 1), 8
        ), reinterpret_tensor(buf0, (4, 4, 4), (64, 16, 1), 12
        ), reinterpret_tensor(buf1, (64, 128), (128, 1), 0), buf4, primals_4


def get_root_logger(log_file=None, log_level=logging.INFO):
    """Use `get_logger` method in mmcv to get the root logger.

    The logger will be initialized if it has not been initialized. By default a
    StreamHandler will be added. If `log_file` is specified, a FileHandler will
    also be added. The name of the root logger is the top-level package name,
    e.g., "mmpose".

    Args:
        log_file (str | None): The log filename. If specified, a FileHandler
            will be added to the root logger.
        log_level (int): The root logger level. Note that only the process of
            rank 0 is affected, while other processes will set the level to
            "Error" and be silent most of the time.

    Returns:
        logging.Logger: The root logger.
    """
    return get_logger(__name__.split('.')[0], log_file, log_level)


class PositionEmbedding2DNew(nn.Module):
    """2D Postion Embedding layer. """

    def __init__(self, max_position_embeddings=128, embedding_dim=128,
        width_embedding=False, height_embedding=False):
        """
        Args:
            max_position_embeddings (int): max normalized input dimension (similar to vocab_size).
            embedding_dim (int): size of embedding vector.
            width_embedding (bool): whether to include width embedding.
            height_embedding (bool): whether to include height embedding.
        """
        super().__init__()
        self.max_position_embeddings = max_position_embeddings
        self.pos_embedding_dim = embedding_dim
        self.x_embedding = nn.Embedding(self.max_position_embeddings, self.
            pos_embedding_dim)
        self.y_embedding = nn.Embedding(self.max_position_embeddings, self.
            pos_embedding_dim)
        self.width_embedding = None
        if width_embedding:
            self.width_embedding = nn.Embedding(self.
                max_position_embeddings, self.pos_embedding_dim)
        self.height_embedding = None
        if height_embedding:
            self.height_embedding = nn.Embedding(self.
                max_position_embeddings, self.pos_embedding_dim)
        self.pos_input_proj = nn.Linear(self.pos_embedding_dim, self.
            pos_embedding_dim)
        self.pos_input_proj_relu = nn.ReLU()

    def init_weights(self, pretrained=None):
        """ Weight initialization

        Args:
            pretrained (str, optional): Path to pre-trained weights. Defaults to None.
        """
        if isinstance(pretrained, str):
            logger = get_root_logger()
            logger.info('Position Embedding:')
            load_checkpoint(self, pretrained, strict=False, logger=logger)
        elif pretrained is None:
            return
        else:
            raise TypeError('pretrained must be a str or None')

    @property
    def with_width_embedding(self):
        """

        Returns:
            Determine the model with the width_embedding or not
        """
        return hasattr(self, 'width_embedding'
            ) and self.width_embedding is not None

    @property
    def with_height_embedding(self):
        """

        Returns:
            Determine the model with the height_embedding or not
        """
        return hasattr(self, 'height_embedding'
            ) and self.height_embedding is not None

    def forward(self, input_0):
        primals_2 = self.x_embedding.weight
        primals_3 = self.y_embedding.weight
        primals_4 = self.pos_input_proj.weight
        primals_5 = self.pos_input_proj.bias
        primals_1 = input_0
        output = call([primals_1, primals_2, primals_3, primals_4, primals_5])
        return output[0]


Measured latency_ms: 0.055398