submission 512883

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


@triton.jit
def reduce_kernel(
    x_ptr,
    output_ptr,
    N,
    BLOCK_SIZE: tl.constexpr,
    NUM_BLOCKS: tl.constexpr,
):
    pid = tl.program_id(0)
    acc = tl.zeros([BLOCK_SIZE], dtype=tl.float64)
    num_chunks = tl.cdiv(N, BLOCK_SIZE)

    chunk = pid
    while chunk < num_chunks:
        offsets = chunk * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
        mask = offsets < N
        x = tl.load(x_ptr + offsets, mask=mask, other=0.0).to(tl.float64)
        acc = acc + x
        chunk += NUM_BLOCKS

    local_sum = tl.sum(acc, axis=0).to(tl.float32)
    tl.atomic_add(output_ptr, local_sum)


def custom_kernel(data: input_t) -> output_t:
    x, output = data
    N = x.numel()

    if N <= 4096:
        output[0] = x.to(torch.float64).sum().to(torch.float32)
        return output[0]

    output.zero_()

    NUM_BLOCKS = 512
    BLOCK_SIZE = 4096

    reduce_kernel[(NUM_BLOCKS,)](
        x, output, N,
        BLOCK_SIZE=BLOCK_SIZE,
        NUM_BLOCKS=NUM_BLOCKS,
        num_warps=16,
    )

    return output[0]