submission 612483

import os
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"

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

cuda_source = """
#include <torch/extension.h>
#include <cub/cub.cuh>

// Persistent temp storage to avoid re-allocation
static void* g_temp = nullptr;
static size_t g_temp_bytes = 0;

torch::Tensor cub_radix_sort(torch::Tensor input, torch::Tensor output) {
	const int n = input.numel();
	float* d_in = input.data_ptr<float>();
	float* d_out = output.data_ptr<float>();

	// Determine temp storage needed
	void* d_temp = nullptr;
	size_t temp_bytes = 0;
	cub::DeviceRadixSort::SortKeys(d_temp, temp_bytes, d_in, d_out, n);

	// Re-allocate only if needed
	if (temp_bytes > g_temp_bytes) {
		if (g_temp) cudaFree(g_temp);
		cudaMalloc(&g_temp, temp_bytes);
		g_temp_bytes = temp_bytes;
	}

	cub::DeviceRadixSort::SortKeys(g_temp, g_temp_bytes, d_in, d_out, n);

	return output;
}
"""

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

module = load_inline(
	name="cub_radix_sort",
	cpp_sources=[cpp_source],
	cuda_sources=[cuda_source],
	functions=["cub_radix_sort"],
	extra_cuda_cflags=["-O3", "--use_fast_math", "-std=c++17"],
	verbose=False,
)


def custom_kernel(data: input_t) -> output_t:
	inp, out = data
	module.cub_radix_sort(inp, out)
	return out