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
from torch import 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_add_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 = 0.0
    tmp2 = tmp0 + tmp1
    tl.store(out_ptr0 + x0, tmp2, xmask)


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


def conv_bn_relu(C_in, C_out, kernel_size, stride, padding, affine=True):
    return nn.Sequential(nn.Conv2d(C_in, C_out, kernel_size, stride=stride,
        padding=padding, bias=False), nn.BatchNorm2d(C_out, affine=affine),
        nn.ReLU(inplace=False))


class AggregateCell(nn.Module):
    """
        before aggregating, both paths should undergo
        conv1x1 with the output channels being equal to the smallest channel size among them
        upsampled to the highest resolution among them
        pre_transform: whether to do convbnrelu before summing up
    """

    def __init__(self, size_1, size_2, agg_size, pre_transform=True):
        super(AggregateCell, self).__init__()
        self.pre_transform = pre_transform
        if self.pre_transform:
            self.branch_1 = conv_bn_relu(size_1, agg_size, 1, 1, 0)
            self.branch_2 = conv_bn_relu(size_2, agg_size, 1, 1, 0)

    def forward(self, x1, x2):
        if self.pre_transform:
            x1 = self.branch_1(x1)
            x2 = self.branch_2(x2)
        if x1.size()[2:] > x2.size()[2:]:
            x2 = nn.Upsample(size=x1.size()[2:], mode='bilinear')(x2)
        elif x1.size()[2:] < x2.size()[2:]:
            x1 = nn.Upsample(size=x2.size()[2:], mode='bilinear')(x1)
        return x1 + x2


class ContextualCellNew(nn.Module):
    """New contextual cell design

    Config contains [op1, [loc1, loc2, op1, op2], [...], [...]]

    """

    def __init__(self, config, inp, repeats=1):
        super(ContextualCellNew, self).__init__()
        self._ops = nn.ModuleList()
        self._pos = []
        self._collect_inds = [0]
        self._pools = ['x']
        for ind, op in enumerate(config):
            if ind == 0:
                pos = 0
                op_id = op
                self._collect_inds.remove(pos)
                op_name = OP_NAMES[op_id]
                self._ops.append(OPS[op_name](inp, inp, 1, True, repeats))
                self._pos.append(pos)
                self._collect_inds.append(ind + 1)
                self._pools.append('{}({})'.format(op_name, self._pools[pos]))
            else:
                pos1, pos2, op_id1, op_id2 = op
                for pos, op_id in zip([pos1, pos2], [op_id1, op_id2]):
                    if pos in self._collect_inds:
                        self._collect_inds.remove(pos)
                    op_name = OP_NAMES[op_id]
                    self._ops.append(OPS[op_name](inp, inp, 1, True, repeats))
                    self._pos.append(pos)
                    self._pools.append('{}({})'.format(op_name, self._pools
                        [pos]))
                op_name = 'sum'
                self._ops.append(AggregateCell(size_1=None, size_2=None,
                    agg_size=inp, pre_transform=False))
                self._pos.append([ind * 3 - 1, ind * 3])
                self._collect_inds.append(ind * 3 + 1)
                self._pools.append('{}({},{})'.format(op_name, self._pools[
                    ind * 3 - 1], self._pools[ind * 3]))

    def prettify(self):
        return ' + '.join(self._pools[i] for i in self._collect_inds)

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


Measured latency_ms: 0.014119