submission 782547

#!POPCORN leaderboard vectorsum_v2

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

BLOCK_SIZE = 16384
NUM_WARPS = 2


@triton.jit
def sum_kernel(input_ptr, partials_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
    pid = tl.program_id(0)
    offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
    mask = offsets < n_elements
    x = tl.load(input_ptr + offsets, mask=mask, other=0.0)
    x = x.to(tl.float64)
    block_sum = tl.sum(x, axis=0)
    tl.store(partials_ptr + pid, block_sum.to(tl.float32))


def custom_kernel(data: input_t) -> output_t:
    input_tensor, _ = data
    n_elements = input_tensor.numel()
    num_blocks = triton.cdiv(n_elements, BLOCK_SIZE)
    partials = torch.empty(num_blocks, device="cuda", dtype=torch.float32)
    sum_kernel[(num_blocks,)](
        input_tensor, partials, n_elements,
        BLOCK_SIZE=BLOCK_SIZE, num_warps=NUM_WARPS,
    )
    return partials.to(torch.float64).sum().to(torch.float32)