submission 132603

#!POPCORN leaderboard prefixsum_v2

# This is a submission template for popcorn leaderboard 'prefixsum_v2'.
# Your task is as follows:
# > Implement an inclusive prefix sum (scan) kernel that matches the reference implementation.
# > The kernel should compute the cumulative sum of all elements up to each position.
# > Because of numerical instability, the tolerance is scaled by the square root of the input size.
# > 
# > Input:
# > - `data`: A 1D tensor of size `n`
# > Output:
# > - `output`: A 1D tensor of size `n`
# The deadline for this leaderboard is 2025-12-30 00:00:00+00:00

# You can automatically route this file to specific GPUs by adding a line
# `#!POPCORN gpus <GPUs>` to the header of this file.
# Happy hacking!

import subprocess

subprocess.run(["pip", "install", "cuda-cccl[cu12]==0.4.1"])

from task import input_t, output_t

import cuda.compute
from cuda.compute import OpKind
import torch

build_in = torch.empty(1, dtype=torch.float32, device="cuda")
build_out = torch.empty(1, dtype=torch.float32, device="cuda")

scanner = cuda.compute.make_inclusive_scan(build_in, build_out, OpKind.PLUS, None)

input_size = 268435456
temp_storage_size = scanner(None, build_in, build_out, input_size, None)
d_temp_storage = torch.empty(temp_storage_size, dtype=torch.uint8, device="cuda")

def custom_kernel(data: input_t) -> output_t:
    d_in, d_out = data
    scanner(d_temp_storage, d_in, d_out, len(d_in), None)


    return d_out