submission 122955

import math
from typing import Optional

import torch

try:
    import triton
    import triton.language as tl

    _HAS_TRITON = True
except Exception:
    _HAS_TRITON = False

from task import input_t, output_t


if _HAS_TRITON:

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


def _triton_reduce(x: torch.Tensor) -> Optional[torch.Tensor]:
    """Fast two-pass reduction using Triton; returns None if Triton unavailable."""
    if not _HAS_TRITON:
        return None
    if not x.is_cuda:
        return None

    # Ensure contiguous for coalesced loads
    x = x.contiguous()
    n = x.numel()
    if n == 0:
        return torch.zeros((), device=x.device, dtype=torch.float32)

    BLOCK = 2048 if n >= (1 << 20) else 1024
    num_warps = 8 if BLOCK >= 2048 else 4
    grid = (triton.cdiv(n, BLOCK),)

    partial = torch.empty(grid[0], device=x.device, dtype=torch.float64)

    _block_sum_kernel[grid](x, partial, n, BLOCK=BLOCK, num_warps=num_warps, num_stages=2)

    # Final reduction with PyTorch on a small buffer; accumulate in float64.
    return torch.sum(partial, dtype=torch.float64).to(torch.float32)


def custom_kernel(data: input_t) -> output_t:
    """
    Sum all elements of the input vector. Accumulates in float64 to match the
    reference, then returns a 0-d float32 tensor. Uses a fast Triton two-pass
    reduction on CUDA; falls back to torch.sum otherwise.
    """
    x, out = data

    total = _triton_reduce(x)
    if total is None:
        total = torch.sum(x, dtype=torch.float64).to(torch.float32)

    try:
        out.view(-1)[0] = total
    except Exception:
        pass

    return total