submission 763415

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

sort_cuda_source = """
#include <cub/cub.cuh>
#include <cuda_runtime.h>

static void* d_temp = nullptr;
static size_t temp_bytes = 0;

void sort_cuda(torch::Tensor input, torch::Tensor output) {
    int N = input.numel();
    float* d_in = input.data_ptr<float>();
    float* d_out = output.data_ptr<float>();

    size_t required = 0;
    cub::DeviceRadixSort::SortKeys(nullptr, required, d_in, d_out, N);

    if (required > temp_bytes) {
        if (d_temp) cudaFree(d_temp);
        cudaMalloc(&d_temp, required);
        temp_bytes = required;
    }

    cub::DeviceRadixSort::SortKeys(d_temp, temp_bytes, d_in, d_out, N);
}
"""

sort_cpp_source = """
#include <torch/extension.h>
void sort_cuda(torch::Tensor input, torch::Tensor output);
"""

sort_module = load_inline(
    name='sort_cuda',
    cpp_sources=sort_cpp_source,
    cuda_sources=sort_cuda_source,
    functions=['sort_cuda'],
    verbose=True,
    extra_cuda_cflags=['-O3', '--use_fast_math', '-gencode', 'arch=compute_90,code=sm_90'],
)

def custom_kernel(data: input_t) -> output_t:
    data, output = data
    sort_module.sort_cuda(data, output)
    return output