import torch
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 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


@triton.jit
def triton_poi_fused_cat_0(in_ptr0, in_ptr1, out_ptr0, xnumel, XBLOCK: tl.
    constexpr):
    xnumel = 512
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 16 % 8
    x0 = xindex % 16
    x2 = xindex // 128
    x3 = xindex
    tmp0 = x1
    tl.full([1], 0, tl.int64)
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = tl.load(in_ptr0 + (x0 + 16 * x1 + 64 * x2), tmp4 & xmask, other=0.0)
    tmp6 = tmp0 >= tmp3
    tl.full([1], 8, tl.int64)
    tmp9 = tl.load(in_ptr1 + (x0 + 16 * (-4 + x1) + 64 * x2), tmp6 & xmask,
        other=0.0)
    tmp10 = tl.where(tmp4, tmp5, tmp9)
    tl.store(out_ptr0 + x3, tmp10, xmask)


def call(args):
    arg0_1, arg1_1 = args
    args.clear()
    assert_size_stride(arg0_1, (4, 4, 4, 4), (64, 16, 4, 1))
    assert_size_stride(arg1_1, (4, 4, 4, 4), (64, 16, 4, 1))
    with torch.cuda._DeviceGuard(0):
        torch.cuda.set_device(0)
        buf0 = empty_strided_cuda((4, 8, 4, 4), (128, 16, 4, 1), torch.float32)
        get_raw_stream(0)
        triton_poi_fused_cat_0[grid(512)](arg0_1, arg1_1, buf0, 512, XBLOCK
            =256, num_warps=4, num_stages=1)
        del arg0_1
        del arg1_1
    return buf0,


class ConcatNew(nn.Module):

    def __init__(self, args=None):
        super(ConcatNew, self).__init__()
        self.index = -1
        self.verbose = print
        self.enable = False
        self.input_index = ''
        self.tag = 'fm'
        self.args = args
        if self.args is not None:
            logger_root = args.logger_root + '.' if hasattr(args, 'logger_root'
                ) else ''
            self.logger = logging.getLogger(logger_root + __name__)
        else:
            self.logger = logging.getLogger(__name__ + '.Quantization')
        self.verbose = self.logger.info

    def update_concat_quantization_parameter(self, **parameters):
        index = self.index
        if 'index' in parameters:
            index = parameters['index']
        if index != self.index:
            self.index = index
            self.verbose('update %s_index %r' % (self.tag, self.index))
        if 'by_index' in parameters:
            by_index = parameters['by_index']
            if isinstance(by_index, list) or isinstance(by_index, str
                ) and by_index != 'all':
                try:
                    if not isinstance(by_index, list):
                        by_index = by_index.split()
                    by_index = [int(i) for i in by_index]
                except (ValueError, SyntaxError):
                    self.verbose('unexpect string in by_index: {}'.format(
                        by_index))
            if by_index == 'all' or self.index in by_index:
                if 'by_tag' in parameters and self.tag in parameters['by_tag'
                    ] or 'by_tag' not in parameters:
                    for k, v in list(parameters.items()):
                        if hasattr(self, '{}'.format(k)):
                            if isinstance(v, str):
                                v = v.replace("'", '').replace('"', '')
                            if isinstance(getattr(self, k), bool):
                                v = False if v in ['False', 'false', False
                                    ] else True
                            elif isinstance(getattr(self, k), int):
                                v = int(v)
                            if not isinstance(getattr(self, k), torch.Tensor):
                                setattr(self, '{}'.format(k), v)
                                self.verbose('update {}_{} to {} for index {}'
                                    .format(self.tag, k, getattr(self, k,
                                    'Non-Exist'), self.index))

    def __repr__(self):
        base = super(ConcatNew, self).__repr__()
        if self.enable:
            base = base + '-index({})-input({})'.format(self.index, self.
                input_index)
        return base

    def forward(self, input_0, input_1):
        arg0_1 = input_0
        arg1_1 = input_1
        output = call([arg0_1, arg1_1])
        return output[0]


Measured latency_ms: 0.015457