submission 612908

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 = r"""
#include <torch/extension.h>
#include <cub/cub.cuh>

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();
	int* d_in = reinterpret_cast<int*>(input.data_ptr<float>());
	int* d_out = reinterpret_cast<int*>(output.data_ptr<float>());

	void* d_temp = nullptr;
	size_t temp_bytes = 0;
	cub::DeviceRadixSort::SortKeys(d_temp, temp_bytes, d_in, d_out, n);

	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_fast",
	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