submission 569099

import torch
import triton
import triton.language as tl
from task import input_t, output_t  # evaluator types

@triton.jit
def _vecsum_kernel(X, PARTIAL, N, BLOCK_SIZE: tl.constexpr):
    """
    Each program sums a block of BLOCK_SIZE elements from X
    and writes the partial sum to PARTIAL.
    """
    pid = tl.program_id(0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < N

    # Load input elements
    x = tl.load(X + offsets, mask=mask, other=0.0)

    # Accumulate in float64 for accuracy
    x = x.to(tl.float64)
    s = tl.sum(x, axis=0)

    # Store partial sum
    tl.store(PARTIAL + pid, s)


def custom_kernel(data: input_t) -> output_t:
    """
    Fast vector sum reduction using Triton.
    Produces identical output as reference (float64 accumulation → float32 scalar).
    """
    data, _ = data
    N = data.numel()
    BLOCK_SIZE = 1024

    # Compute number of programs
    num_blocks = triton.cdiv(N, BLOCK_SIZE)

    # Allocate partial sums
    partial = torch.empty(num_blocks, device=data.device, dtype=torch.float64)

    # Launch kernel
    _vecsum_kernel[(num_blocks,)](data, partial, N, BLOCK_SIZE=BLOCK_SIZE)

    # Final reduction (small tensor, PyTorch is fine)
    result = partial.sum().to(torch.float32)

    return result