submission 150650

#!POPCORN leaderboard sort_v2

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

# CUB-based radix sort implementation
cuda_source = """
#include <torch/extension.h>
#include <cub/device/device_radix_sort.cuh>
#include <cuda_runtime.h>
#include <c10/cuda/CUDACachingAllocator.h>

void cub_radix_sort_float(
    const float* keys_in,
    float* keys_out,
    int64_t n) {

    // Determine temporary device storage requirements
    void* d_temp_storage = nullptr;
    size_t temp_storage_bytes = 0;

    cub::DeviceRadixSort::SortKeys(
        d_temp_storage, temp_storage_bytes,
        keys_in, keys_out, n);

    // Allocate temporary storage using PyTorch's caching allocator
    auto& allocator = *c10::cuda::CUDACachingAllocator::get();
    auto temp_storage = allocator.allocate(temp_storage_bytes);
    d_temp_storage = temp_storage.get();

    // Run sorting operation
    cub::DeviceRadixSort::SortKeys(
        d_temp_storage, temp_storage_bytes,
        keys_in, keys_out, n);
    // temp_storage is automatically freed when it goes out of scope
}

torch::Tensor sort_kernel_cuda(torch::Tensor input, torch::Tensor output) {
    TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
    TORCH_CHECK(output.is_cuda(), "output must be a CUDA tensor");
    TORCH_CHECK(input.is_contiguous(), "input must be contiguous");
    TORCH_CHECK(output.is_contiguous(), "output must be contiguous");
    TORCH_CHECK(input.dtype() == torch::kFloat32, "input must be float32");

    int64_t n = input.numel();

    const float* input_ptr = input.data_ptr<float>();
    float* output_ptr = output.data_ptr<float>();

    cub_radix_sort_float(input_ptr, output_ptr, n);

    return output;
}
"""

cpp_source = """
torch::Tensor sort_kernel_cuda(torch::Tensor input, torch::Tensor output);
"""

# Compile the CUDA extension
cub_sort_module = load_inline(
    name='cub_sort',
    cpp_sources=cpp_source,
    cuda_sources=cuda_source,
    functions=['sort_kernel_cuda'],
    with_cuda=True,
    extra_cuda_cflags=['-O3', '-use_fast_math'],
)

def cub_kernel(data: input_t) -> output_t:
    inp, output = data
    cub_sort_module.sort_kernel_cuda(inp, output)
    return output

def naive_custom_kernel(data):
    inp, out = data

    out.copy_(torch.sort(inp)[0])
    return out


custom_kernel = cub_kernel