submission 132576

#!POPCORN leaderboard sort_v2

# This is a submission template for popcorn leaderboard 'sort_v2'.
# Your task is as follows:
# > Implement a sort kernel that matches the reference implementation.
# > The kernel should sort the input array in ascending order using a sort algorithm of your choice.
# > 
# > Input arrays are generated as random floating-point numbers, where each row of a roughly square matrix
# > is drawn from a normal distribution with a different mean value per row based on the seed and then flattened into a 1D array.
# 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 (
    DoubleBuffer,
    SortOrder,
)
import torch

build_in_keys = torch.empty(1, dtype=torch.float32, device="cuda")
build_out_keys = torch.empty(1, dtype=torch.float32, device="cuda")

build_double_buffer = DoubleBuffer(build_in_keys, build_out_keys)

sorter = cuda.compute.make_radix_sort(build_double_buffer, None, None, None, SortOrder.ASCENDING)

input_size = 100000000

temp_storage_size = sorter(None, build_double_buffer, None, None, None, input_size)
d_temp_storage = torch.empty(temp_storage_size, dtype=torch.uint8, device="cuda")

def custom_kernel(data: input_t) -> output_t:
    d_in_keys, d_out_keys = data
    double_buffer = DoubleBuffer(d_in_keys, d_out_keys)
    size = len(d_in_keys)

    sorter(d_temp_storage, double_buffer, None, None, None, size)

    return double_buffer.current()