submission 768762

#!POPCORN leaderboard prefixsum_v2
import torch
from task import input_t, output_t

CUDA_SOURCE = r"""
// naive serial prefix sum — one thread does all the work
__global__ void solution(const float* __restrict__ input, float* __restrict__ output, int n) {
    if (threadIdx.x == 0 && blockIdx.x == 0) {
        double acc = 0.0;
        for (int i = 0; i < n; i++) {
            acc += (double)input[i];
            output[i] = (float)acc;
        }
    }
}

"""

_kernel = torch.cuda._compile_kernel(CUDA_SOURCE, "solution")

def custom_kernel(data: input_t) -> output_t:
    x, y = data
    n = x.numel()
    # adjust grid/block to match your kernel
    threads = (256, 1, 1)
    blocks = ((n + 255) // 256, 1, 1)
    _kernel(blocks, threads, (x, y, n))
    return y