submission 682469

#!POPCORN leaderboard sort_v2
#!POPCORN gpu B200

import torch
from torch.utils.cpp_extension import load_inline
from task import input_t, output_t

cuda_src = r"""
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cub/cub.cuh>

// Pre-allocated temp storage
static void* d_temp = nullptr;
static size_t d_temp_size = 0;

void sort_init(int64_t max_n) {
    // Query temp storage size for max problem size
    size_t needed = 0;
    cub::DeviceRadixSort::SortKeys(
        nullptr, needed,
        (const float*)nullptr, (float*)nullptr,
        (int)max_n);

    if (needed > d_temp_size) {
        if (d_temp) cudaFree(d_temp);
        // Allocate with some headroom
        d_temp_size = needed * 2;
        cudaMalloc(&d_temp, d_temp_size);
    }
}

void sort_float(int64_t in_ptr, int64_t out_ptr, int N) {
    size_t temp_size = d_temp_size;
    cub::DeviceRadixSort::SortKeys(
        d_temp, temp_size,
        reinterpret_cast<const float*>(in_ptr),
        reinterpret_cast<float*>(out_ptr),
        N);
}
"""

cpp_src = r"""
void sort_init(int64_t max_n);
void sort_float(int64_t in_ptr, int64_t out_ptr, int N);
"""

_ext = load_inline(
    name="cub_radix_sort",
    cpp_sources=cpp_src,
    cuda_sources=cuda_src,
    functions=["sort_init", "sort_float"],
    with_cuda=True,
    extra_cflags=["-O3"],
    extra_cuda_cflags=["-O3", "--use_fast_math", "-arch=sm_100a"],
    verbose=False,
)

# Pre-allocate temp storage for max size (100M elements)
_ext.sort_init(100_000_000)

# Warmup
_wd = torch.randn(1000, device="cuda", dtype=torch.float32)
_wo = torch.empty_like(_wd)
_ext.sort_float(_wd.data_ptr(), _wo.data_ptr(), _wd.numel())
torch.cuda.synchronize()
del _wd, _wo


def custom_kernel(data: input_t) -> output_t:
    data_in, output = data
    _ext.sort_float(data_in.data_ptr(), output.data_ptr(), data_in.numel())
    return output