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 torch.nn as nn
import torch._C
import torch.serialization
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_per_fused_cat_mean_0(in_ptr0, out_ptr1, xnumel, rnumel, XBLOCK:
    tl.constexpr):
    xnumel = 16
    RBLOCK: tl.constexpr = 16
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:, None]
    xmask = xindex < xnumel
    rindex = tl.arange(0, RBLOCK)[None, :]
    tl.full([XBLOCK, RBLOCK], True, tl.int1)
    r1 = rindex
    x0 = xindex
    tmp0 = tl.load(in_ptr0 + (r1 + 16 * x0), xmask, other=0.0)
    tmp1 = tl.broadcast_to(tmp0, [XBLOCK, RBLOCK])
    tmp3 = tl.where(xmask, tmp1, 0)
    tmp4 = tl.sum(tmp3, 1)[:, None]
    tmp5 = 16.0
    tmp6 = tmp4 / tmp5
    tl.store(out_ptr1 + 110 * x0, tmp6, xmask)


@triton.jit
def triton_poi_fused__adaptive_avg_pool2d_cat_1(in_ptr0, out_ptr1, xnumel,
    XBLOCK: tl.constexpr):
    xnumel = 144
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 3 % 3
    x0 = xindex % 3
    x2 = xindex // 9
    x3 = xindex % 9
    tmp0 = 4 * x1 // 3
    tmp1 = 2 + 4 * x1 // 3
    tmp2 = tmp0 < tmp1
    tmp3 = 4 * x0 // 3
    tmp4 = 2 + 4 * x0 // 3
    tmp5 = tmp3 < tmp4
    tmp6 = tmp2 & tmp5
    tmp7 = tl.load(in_ptr0 + (4 * (4 * x1 // 3) + 16 * x2 + 4 * x0 // 3), 
        tmp6 & xmask, other=0.0)
    tmp8 = 1 + 4 * x0 // 3
    tmp9 = tmp8 < tmp4
    tmp10 = tmp2 & tmp9
    tmp11 = tl.load(in_ptr0 + (1 + 4 * (4 * x1 // 3) + 16 * x2 + 4 * x0 // 
        3), tmp10 & xmask, other=0.0)
    tmp12 = tmp11 + tmp7
    tmp13 = 1 + 4 * x1 // 3
    tmp14 = tmp13 < tmp1
    tmp15 = tmp14 & tmp5
    tmp16 = tl.load(in_ptr0 + (4 + 4 * (4 * x1 // 3) + 16 * x2 + 4 * x0 // 
        3), tmp15 & xmask, other=0.0)
    tmp17 = tmp16 + tmp12
    tmp18 = tmp14 & tmp9
    tmp19 = tl.load(in_ptr0 + (5 + 4 * (4 * x1 // 3) + 16 * x2 + 4 * x0 // 
        3), tmp18 & xmask, other=0.0)
    tmp20 = tmp19 + tmp17
    tmp21 = 1.0
    tmp22 = tl.full(tmp21.shape, 0.0, tmp21.dtype)
    tmp23 = tl.where(tmp6, tmp21, tmp22)
    tmp24 = tl.where(tmp10, tmp21, tmp22)
    tmp25 = tmp24 + tmp23
    tmp26 = tl.where(tmp15, tmp21, tmp22)
    tmp27 = tmp26 + tmp25
    tmp28 = tl.where(tmp18, tmp21, tmp22)
    tmp29 = tmp28 + tmp27
    tmp30 = tmp20 / tmp29
    tl.store(out_ptr1 + (x3 + 110 * x2), tmp30, xmask)


@triton.jit
def triton_poi_fused__adaptive_avg_pool2d_cat_2(in_ptr0, out_ptr1, xnumel,
    XBLOCK: tl.constexpr):
    xnumel = 576
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 6 % 6
    x0 = xindex % 6
    x2 = xindex // 36
    x3 = xindex % 36
    tmp0 = 2 * x1 // 3
    tmp1 = (9 + 4 * x1) // 6
    tmp2 = tmp0 < tmp1
    tmp3 = 2 * x0 // 3
    tmp4 = (9 + 4 * x0) // 6
    tmp5 = tmp3 < tmp4
    tmp6 = tmp2 & tmp5
    tmp7 = tl.load(in_ptr0 + (4 * (2 * x1 // 3) + 16 * x2 + 2 * x0 // 3), 
        tmp6 & xmask, eviction_policy='evict_last', other=0.0)
    tmp8 = 1 + 2 * x0 // 3
    tmp9 = tmp8 < tmp4
    tmp10 = tmp2 & tmp9
    tmp11 = tl.load(in_ptr0 + (1 + 4 * (2 * x1 // 3) + 16 * x2 + 2 * x0 // 
        3), tmp10 & xmask, eviction_policy='evict_last', other=0.0)
    tmp12 = tmp11 + tmp7
    tmp13 = 1 + 2 * x1 // 3
    tmp14 = tmp13 < tmp1
    tmp15 = tmp14 & tmp5
    tmp16 = tl.load(in_ptr0 + (4 + 4 * (2 * x1 // 3) + 16 * x2 + 2 * x0 // 
        3), tmp15 & xmask, eviction_policy='evict_last', other=0.0)
    tmp17 = tmp16 + tmp12
    tmp18 = tmp14 & tmp9
    tmp19 = tl.load(in_ptr0 + (5 + 4 * (2 * x1 // 3) + 16 * x2 + 2 * x0 // 
        3), tmp18 & xmask, eviction_policy='evict_last', other=0.0)
    tmp20 = tmp19 + tmp17
    tmp21 = 1.0
    tmp22 = tl.full(tmp21.shape, 0.0, tmp21.dtype)
    tmp23 = tl.where(tmp6, tmp21, tmp22)
    tmp24 = tl.where(tmp10, tmp21, tmp22)
    tmp25 = tmp24 + tmp23
    tmp26 = tl.where(tmp15, tmp21, tmp22)
    tmp27 = tmp26 + tmp25
    tmp28 = tl.where(tmp18, tmp21, tmp22)
    tmp29 = tmp28 + tmp27
    tmp30 = tmp20 / tmp29
    tl.store(out_ptr1 + (x3 + 110 * x2), tmp30, xmask)


@triton.jit
def triton_poi_fused__adaptive_avg_pool2d_cat_3(in_ptr0, out_ptr1, xnumel,
    XBLOCK: tl.constexpr):
    xnumel = 1024
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 8 % 8
    x0 = xindex % 8
    x2 = xindex // 64
    x3 = xindex % 64
    tmp0 = x1 // 2
    tmp1 = (11 + 4 * x1) // 8
    tmp2 = tmp0 < tmp1
    tmp3 = x0 // 2
    tmp4 = (11 + 4 * x0) // 8
    tmp5 = tmp3 < tmp4
    tmp6 = tmp2 & tmp5
    tmp7 = tl.load(in_ptr0 + (4 * (x1 // 2) + 16 * x2 + x0 // 2), tmp6 &
        xmask, eviction_policy='evict_last', other=0.0)
    tmp8 = 1 + x0 // 2
    tmp9 = tmp8 < tmp4
    tmp10 = tmp2 & tmp9
    tmp11 = tl.load(in_ptr0 + (1 + 4 * (x1 // 2) + 16 * x2 + x0 // 2), 
        tmp10 & xmask, eviction_policy='evict_last', other=0.0)
    tmp12 = tmp11 + tmp7
    tmp13 = 1 + x1 // 2
    tmp14 = tmp13 < tmp1
    tmp15 = tmp14 & tmp5
    tmp16 = tl.load(in_ptr0 + (4 + 4 * (x1 // 2) + 16 * x2 + x0 // 2), 
        tmp15 & xmask, eviction_policy='evict_last', other=0.0)
    tmp17 = tmp16 + tmp12
    tmp18 = tmp14 & tmp9
    tmp19 = tl.load(in_ptr0 + (5 + 4 * (x1 // 2) + 16 * x2 + x0 // 2), 
        tmp18 & xmask, eviction_policy='evict_last', other=0.0)
    tmp20 = tmp19 + tmp17
    tmp21 = 1.0
    tmp22 = tl.full(tmp21.shape, 0.0, tmp21.dtype)
    tmp23 = tl.where(tmp6, tmp21, tmp22)
    tmp24 = tl.where(tmp10, tmp21, tmp22)
    tmp25 = tmp24 + tmp23
    tmp26 = tl.where(tmp15, tmp21, tmp22)
    tmp27 = tmp26 + tmp25
    tmp28 = tl.where(tmp18, tmp21, tmp22)
    tmp29 = tmp28 + tmp27
    tmp30 = tmp20 / tmp29
    tl.store(out_ptr1 + (x3 + 110 * x2), tmp30, 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)
        buf8 = empty_strided_cuda((4, 4, 110), (440, 110, 1), torch.float32)
        buf4 = reinterpret_tensor(buf8, (4, 4, 1), (440, 110, 1), 0)
        get_raw_stream(0)
        triton_per_fused_cat_mean_0[grid(16)](arg0_1, buf4, 16, 16, XBLOCK=
            1, num_warps=2, num_stages=1)
        buf5 = reinterpret_tensor(buf8, (4, 4, 9), (440, 110, 1), 1)
        triton_poi_fused__adaptive_avg_pool2d_cat_1[grid(144)](arg0_1, buf5,
            144, XBLOCK=128, num_warps=4, num_stages=1)
        buf6 = reinterpret_tensor(buf8, (4, 4, 36), (440, 110, 1), 10)
        triton_poi_fused__adaptive_avg_pool2d_cat_2[grid(576)](arg0_1, buf6,
            576, XBLOCK=128, num_warps=4, num_stages=1)
        buf7 = reinterpret_tensor(buf8, (4, 4, 64), (440, 110, 1), 46)
        triton_poi_fused__adaptive_avg_pool2d_cat_3[grid(1024)](arg0_1,
            buf7, 1024, XBLOCK=128, num_warps=4, num_stages=1)
        del arg0_1
    return buf8,


class PPMConcatNew(nn.ModuleList):
    """Pyramid Pooling Module that only concat the features of each layer.

    Args:
        pool_scales (tuple[int]): Pooling scales used in Pooling Pyramid
            Module.
    """

    def __init__(self, pool_scales=(1, 3, 6, 8)):
        super(PPMConcatNew, self).__init__([nn.AdaptiveAvgPool2d(pool_scale
            ) for pool_scale in pool_scales])

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


Measured latency_ms: 0.03978