submission 512009

import torch

try:
    import triton
    import triton.language as tl

    _TRITON_AVAILABLE = True
except Exception:
    triton = None
    tl = None
    _TRITON_AVAILABLE = False

from task import input_t, output_t


BLOCK_SIZE = 1024
FIRST_PASS_CHUNK = 2
LARGE_SIZE_THRESHOLD = 1 << 24
FIRST_PASS_WARPS_SMALL = 8
FIRST_PASS_WARPS_LARGE = 4
LATE_PASS_WARPS = 4

_BUF_A = {}
_BUF_B = {}


if _TRITON_AVAILABLE:

    @triton.jit
    def _reduce_block_kernel(
        x_ptr,
        partial_ptr,
        n_elements,
        BLOCK_SIZE: tl.constexpr,
    ):
        pid = tl.program_id(axis=0)
        offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
        mask = offsets < n_elements
        vals = tl.load(x_ptr + offsets, mask=mask, other=0.0).to(tl.float64)
        acc = tl.sum(vals, axis=0)
        tl.store(partial_ptr + pid, acc)

    @triton.jit
    def _reduce_chunk_kernel(
        x_ptr,
        partial_ptr,
        n_elements,
        BLOCK_SIZE: tl.constexpr,
        CHUNK: tl.constexpr,
    ):
        pid = tl.program_id(axis=0)
        base = pid * BLOCK_SIZE * CHUNK
        acc = tl.zeros((), dtype=tl.float64)
        for chunk_idx in tl.static_range(0, CHUNK):
            offsets = base + chunk_idx * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
            mask = offsets < n_elements
            vals = tl.load(x_ptr + offsets, mask=mask, other=0.0).to(tl.float64)
            acc += tl.sum(vals, axis=0)
        tl.store(partial_ptr + pid, acc)


def _get_buffer(cache: dict, device: torch.device, needed: int) -> torch.Tensor:
    key = (device.type, device.index)
    buf = cache.get(key)
    if buf is None or buf.numel() < needed:
        alloc = 1 << (max(1, needed) - 1).bit_length()
        buf = torch.empty((alloc,), device=device, dtype=torch.float64)
        cache[key] = buf
    return buf[:needed]


def _launch_reduce_block(
    inp: torch.Tensor,
    out: torch.Tensor,
    n_elements: int,
    num_warps: int,
) -> None:
    grid = (out.numel(),)
    _reduce_block_kernel[grid](
        inp,
        out,
        n_elements,
        BLOCK_SIZE=BLOCK_SIZE,
        num_warps=num_warps,
        num_stages=2,
    )


def _launch_reduce_chunk_first(
    inp: torch.Tensor,
    out: torch.Tensor,
    n_elements: int,
    num_warps: int,
) -> None:
    grid = (out.numel(),)
    _reduce_chunk_kernel[grid](
        inp,
        out,
        n_elements,
        BLOCK_SIZE=BLOCK_SIZE,
        CHUNK=FIRST_PASS_CHUNK,
        num_warps=num_warps,
        num_stages=2,
    )


def _triton_sum_fp64(x: torch.Tensor) -> torch.Tensor:
    n0 = x.numel()

    if n0 >= LARGE_SIZE_THRESHOLD:
        b1 = triton.cdiv(n0, BLOCK_SIZE * FIRST_PASS_CHUNK)
        p1 = _get_buffer(_BUF_A, x.device, b1)
        _launch_reduce_chunk_first(x, p1, n0, FIRST_PASS_WARPS_LARGE)
    else:
        b1 = triton.cdiv(n0, BLOCK_SIZE)
        p1 = _get_buffer(_BUF_A, x.device, b1)
        _launch_reduce_block(x, p1, n0, FIRST_PASS_WARPS_SMALL)

    if b1 == 1:
        return p1[0]

    b2 = triton.cdiv(b1, BLOCK_SIZE)
    p2 = _get_buffer(_BUF_B, x.device, b2)
    _launch_reduce_block(p1, p2, b1, LATE_PASS_WARPS)
    if b2 == 1:
        return p2[0]

    b3 = triton.cdiv(b2, BLOCK_SIZE)
    p3 = _get_buffer(_BUF_A, x.device, b3)
    _launch_reduce_block(p2, p3, b2, LATE_PASS_WARPS)
    if b3 == 1:
        return p3[0]

    current = p3
    n_current = b3
    use_a = False
    while n_current > 1:
        b = triton.cdiv(n_current, BLOCK_SIZE)
        if use_a:
            nxt = _get_buffer(_BUF_A, x.device, b)
        else:
            nxt = _get_buffer(_BUF_B, x.device, b)
        _launch_reduce_block(current, nxt, n_current, LATE_PASS_WARPS)
        current = nxt
        n_current = b
        use_a = not use_a
    return current[0]


def custom_kernel(data: input_t) -> output_t:
    x, output = data
    if not x.is_contiguous():
        x = x.contiguous()

    if x.numel() == 0:
        total = torch.zeros((), device=x.device, dtype=torch.float64)
    elif x.is_cuda and _TRITON_AVAILABLE:
        try:
            total = _triton_sum_fp64(x)
        except Exception:
            total = x.to(torch.float64).sum()
    else:
        total = x.to(torch.float64).sum()

    out_scalar = output.view(())
    out_scalar.copy_(total)
    return out_scalar