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._inductor.runtime import triton_helpers
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_cat_max_pool2d_with_indices_0(in_ptr0, out_ptr0,
    out_ptr1, xnumel, XBLOCK: tl.constexpr):
    xnumel = 256
    xoffset = tl.program_id(0) * XBLOCK
    xindex = xoffset + tl.arange(0, XBLOCK)[:]
    xmask = xindex < xnumel
    x1 = xindex // 4 % 4
    x0 = xindex % 4
    x7 = xindex
    x3 = xindex // 64
    x4 = xindex % 64
    tmp116 = tl.load(in_ptr0 + x7, xmask)
    tmp0 = -2 + x1
    tmp1 = tl.full([1], 0, tl.int64)
    tmp2 = tmp0 >= tmp1
    tmp3 = tl.full([1], 4, tl.int64)
    tmp4 = tmp0 < tmp3
    tmp5 = tmp2 & tmp4
    tmp6 = -2 + x0
    tmp7 = tmp6 >= tmp1
    tmp8 = tmp6 < tmp3
    tmp9 = tmp7 & tmp8
    tmp10 = tmp5 & tmp9
    tmp11 = tl.load(in_ptr0 + (-10 + x7), tmp10 & xmask, other=float('-inf'))
    tmp12 = -1 + x0
    tmp13 = tmp12 >= tmp1
    tmp14 = tmp12 < tmp3
    tmp15 = tmp13 & tmp14
    tmp16 = tmp5 & tmp15
    tmp17 = tl.load(in_ptr0 + (-9 + x7), tmp16 & xmask, other=float('-inf'))
    tmp18 = triton_helpers.maximum(tmp17, tmp11)
    tmp19 = x0
    tmp20 = tmp19 >= tmp1
    tmp21 = tmp19 < tmp3
    tmp22 = tmp20 & tmp21
    tmp23 = tmp5 & tmp22
    tmp24 = tl.load(in_ptr0 + (-8 + x7), tmp23 & xmask, other=float('-inf'))
    tmp25 = triton_helpers.maximum(tmp24, tmp18)
    tmp26 = 1 + x0
    tmp27 = tmp26 >= tmp1
    tmp28 = tmp26 < tmp3
    tmp29 = tmp27 & tmp28
    tmp30 = tmp5 & tmp29
    tmp31 = tl.load(in_ptr0 + (-7 + x7), tmp30 & xmask, other=float('-inf'))
    tmp32 = triton_helpers.maximum(tmp31, tmp25)
    tmp33 = 2 + x0
    tmp34 = tmp33 >= tmp1
    tmp35 = tmp33 < tmp3
    tmp36 = tmp34 & tmp35
    tmp37 = tmp5 & tmp36
    tmp38 = tl.load(in_ptr0 + (-6 + x7), tmp37 & xmask, other=float('-inf'))
    tmp39 = triton_helpers.maximum(tmp38, tmp32)
    tmp40 = -1 + x1
    tmp41 = tmp40 >= tmp1
    tmp42 = tmp40 < tmp3
    tmp43 = tmp41 & tmp42
    tmp44 = tmp43 & tmp9
    tmp45 = tl.load(in_ptr0 + (-6 + x7), tmp44 & xmask, other=float('-inf'))
    tmp46 = triton_helpers.maximum(tmp45, tmp39)
    tmp47 = tmp43 & tmp15
    tmp48 = tl.load(in_ptr0 + (-5 + x7), tmp47 & xmask, other=float('-inf'))
    tmp49 = triton_helpers.maximum(tmp48, tmp46)
    tmp50 = tmp43 & tmp22
    tmp51 = tl.load(in_ptr0 + (-4 + x7), tmp50 & xmask, other=float('-inf'))
    tmp52 = triton_helpers.maximum(tmp51, tmp49)
    tmp53 = tmp43 & tmp29
    tmp54 = tl.load(in_ptr0 + (-3 + x7), tmp53 & xmask, other=float('-inf'))
    tmp55 = triton_helpers.maximum(tmp54, tmp52)
    tmp56 = tmp43 & tmp36
    tmp57 = tl.load(in_ptr0 + (-2 + x7), tmp56 & xmask, other=float('-inf'))
    tmp58 = triton_helpers.maximum(tmp57, tmp55)
    tmp59 = x1
    tmp60 = tmp59 >= tmp1
    tmp61 = tmp59 < tmp3
    tmp62 = tmp60 & tmp61
    tmp63 = tmp62 & tmp9
    tmp64 = tl.load(in_ptr0 + (-2 + x7), tmp63 & xmask, other=float('-inf'))
    tmp65 = triton_helpers.maximum(tmp64, tmp58)
    tmp66 = tmp62 & tmp15
    tmp67 = tl.load(in_ptr0 + (-1 + x7), tmp66 & xmask, other=float('-inf'))
    tmp68 = triton_helpers.maximum(tmp67, tmp65)
    tmp69 = tmp62 & tmp22
    tmp70 = tl.load(in_ptr0 + x7, tmp69 & xmask, other=float('-inf'))
    tmp71 = triton_helpers.maximum(tmp70, tmp68)
    tmp72 = tmp62 & tmp29
    tmp73 = tl.load(in_ptr0 + (1 + x7), tmp72 & xmask, other=float('-inf'))
    tmp74 = triton_helpers.maximum(tmp73, tmp71)
    tmp75 = tmp62 & tmp36
    tmp76 = tl.load(in_ptr0 + (2 + x7), tmp75 & xmask, other=float('-inf'))
    tmp77 = triton_helpers.maximum(tmp76, tmp74)
    tmp78 = 1 + x1
    tmp79 = tmp78 >= tmp1
    tmp80 = tmp78 < tmp3
    tmp81 = tmp79 & tmp80
    tmp82 = tmp81 & tmp9
    tmp83 = tl.load(in_ptr0 + (2 + x7), tmp82 & xmask, other=float('-inf'))
    tmp84 = triton_helpers.maximum(tmp83, tmp77)
    tmp85 = tmp81 & tmp15
    tmp86 = tl.load(in_ptr0 + (3 + x7), tmp85 & xmask, other=float('-inf'))
    tmp87 = triton_helpers.maximum(tmp86, tmp84)
    tmp88 = tmp81 & tmp22
    tmp89 = tl.load(in_ptr0 + (4 + x7), tmp88 & xmask, other=float('-inf'))
    tmp90 = triton_helpers.maximum(tmp89, tmp87)
    tmp91 = tmp81 & tmp29
    tmp92 = tl.load(in_ptr0 + (5 + x7), tmp91 & xmask, other=float('-inf'))
    tmp93 = triton_helpers.maximum(tmp92, tmp90)
    tmp94 = tmp81 & tmp36
    tmp95 = tl.load(in_ptr0 + (6 + x7), tmp94 & xmask, other=float('-inf'))
    tmp96 = triton_helpers.maximum(tmp95, tmp93)
    tmp97 = 2 + x1
    tmp98 = tmp97 >= tmp1
    tmp99 = tmp97 < tmp3
    tmp100 = tmp98 & tmp99
    tmp101 = tmp100 & tmp9
    tmp102 = tl.load(in_ptr0 + (6 + x7), tmp101 & xmask, other=float('-inf'))
    tmp103 = triton_helpers.maximum(tmp102, tmp96)
    tmp104 = tmp100 & tmp15
    tmp105 = tl.load(in_ptr0 + (7 + x7), tmp104 & xmask, other=float('-inf'))
    tmp106 = triton_helpers.maximum(tmp105, tmp103)
    tmp107 = tmp100 & tmp22
    tmp108 = tl.load(in_ptr0 + (8 + x7), tmp107 & xmask, other=float('-inf'))
    tmp109 = triton_helpers.maximum(tmp108, tmp106)
    tmp110 = tmp100 & tmp29
    tmp111 = tl.load(in_ptr0 + (9 + x7), tmp110 & xmask, other=float('-inf'))
    tmp112 = triton_helpers.maximum(tmp111, tmp109)
    tmp113 = tmp100 & tmp36
    tmp114 = tl.load(in_ptr0 + (10 + x7), tmp113 & xmask, other=float('-inf'))
    tmp115 = triton_helpers.maximum(tmp114, tmp112)
    tl.store(out_ptr0 + (x4 + 256 * x3), tmp115, xmask)
    tl.store(out_ptr1 + (x4 + 256 * x3), tmp116, xmask)


@triton.jit
def triton_poi_fused_cat_1(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
    x2 = xindex
    x0 = xindex % 64
    x1 = xindex // 64
    tmp0 = tl.load(in_ptr0 + x2, xmask)
    tl.store(out_ptr0 + (x0 + 256 * x1), tmp0, 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 = torch.ops.aten.max_pool2d_with_indices.default(arg0_1, [13, 
            13], [1, 1], [6, 6])
        buf1 = buf0[0]
        del buf0
        buf3 = torch.ops.aten.max_pool2d_with_indices.default(arg0_1, [9, 9
            ], [1, 1], [4, 4])
        buf4 = buf3[0]
        del buf3
        buf10 = empty_strided_cuda((4, 16, 4, 4), (256, 16, 4, 1), torch.
            float32)
        buf6 = reinterpret_tensor(buf10, (4, 4, 4, 4), (256, 16, 4, 1), 128)
        buf9 = reinterpret_tensor(buf10, (4, 4, 4, 4), (256, 16, 4, 1), 192)
        get_raw_stream(0)
        triton_poi_fused_cat_max_pool2d_with_indices_0[grid(256)](arg0_1,
            buf6, buf9, 256, XBLOCK=256, num_warps=4, num_stages=1)
        del arg0_1
        buf7 = reinterpret_tensor(buf10, (4, 4, 4, 4), (256, 16, 4, 1), 0)
        triton_poi_fused_cat_1[grid(256)](buf1, buf7, 256, XBLOCK=256,
            num_warps=4, num_stages=1)
        del buf1
        buf8 = reinterpret_tensor(buf10, (4, 4, 4, 4), (256, 16, 4, 1), 64)
        triton_poi_fused_cat_1[grid(256)](buf4, buf8, 256, XBLOCK=256,
            num_warps=4, num_stages=1)
        del buf4
    return buf10,


class SpatialPyramidPoolingNew(nn.Module):

    def __init__(self, pool_sizes=[5, 9, 13]):
        super(SpatialPyramidPoolingNew, self).__init__()
        self.maxpools = nn.ModuleList([nn.MaxPool2d(pool_size, 1, pool_size //
            2) for pool_size in pool_sizes])

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


Measured latency_ms: 0.054757