submission 71047

import torch
from torch.utils.cpp_extension import load_inline

cuda_src = """
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cooperative_groups.h>

namespace cg = cooperative_groups;

#define WARP_SIZE 32
#define RADIX_BITS 8
#define RADIX 256

__device__ __forceinline__ unsigned int float_flip(float f) {
    unsigned int u = __float_as_uint(f);
    return (u & 0x80000000) ? ~u : u ^ 0x80000000;
}

__device__ __forceinline__ float flip_float(unsigned int u) {
    u = (u & 0x80000000) ? u ^ 0x80000000 : ~u;
    return __uint_as_float(u);
}

__global__ void count_radix(const float* __restrict__ in, int* __restrict__ hist,
                            int n, int shift) {
    __shared__ int local[RADIX];
    int tid = threadIdx.x;

    if (tid < RADIX) local[tid] = 0;
    __syncthreads();

    int idx = blockIdx.x * blockDim.x + tid;
    int stride = blockDim.x * gridDim.x;

    for (int i = idx; i < n; i += stride) {
        unsigned int key = float_flip(in[i]);
        int digit = (key >> shift) & (RADIX - 1);
        atomicAdd(&local[digit], 1);
    }
    __syncthreads();

    if (tid < RADIX) {
        atomicAdd(&hist[blockIdx.x * RADIX + tid], local[tid]);
    }
}

__global__ void scan_radix(int* hist, int num_blocks) {
    __shared__ int temp[RADIX];
    int tid = threadIdx.x;

    if (tid < RADIX) {
        int sum = 0;
        for (int b = 0; b < num_blocks; b++) {
            sum += hist[b * RADIX + tid];
        }
        temp[tid] = sum;
    }
    __syncthreads();

    if (tid == 0) {
        int acc = 0;
        for (int i = 0; i < RADIX; i++) {
            int cnt = temp[i];
            temp[i] = acc;
            acc += cnt;
        }
    }
    __syncthreads();

    if (tid < RADIX) hist[tid] = temp[tid];
}

__global__ void reorder_radix(const float* __restrict__ in, float* __restrict__ out,
                              int* __restrict__ offsets, int n, int shift) {
    __shared__ int local_off[RADIX];
    __shared__ int block_off[RADIX];
    int tid = threadIdx.x;

    if (tid < RADIX) local_off[tid] = 0;
    __syncthreads();

    int block_start = blockIdx.x * blockDim.x;
    int block_end = min(block_start + (int)blockDim.x, n);

    for (int i = block_start + tid; i < block_end; i += blockDim.x) {
        unsigned int key = float_flip(in[i]);
        int digit = (key >> shift) & (RADIX - 1);
        atomicAdd(&local_off[digit], 1);
    }
    __syncthreads();

    if (tid < RADIX && local_off[tid] > 0) {
        block_off[tid] = atomicAdd(&offsets[tid], local_off[tid]);
    }
    __syncthreads();

    __shared__ int write_off[RADIX];
    if (tid < RADIX) write_off[tid] = 0;
    __syncthreads();

    for (int i = block_start + tid; i < block_end; i += blockDim.x) {
        float val = in[i];
        unsigned int key = float_flip(val);
        int digit = (key >> shift) & (RADIX - 1);
        int pos = block_off[digit] + atomicAdd(&write_off[digit], 1);
        out[pos] = val;
    }
}

void radix_sort_impl(const float* in, float* out, float* tmp, int n, cudaStream_t stream) {
    const int threads = 1024;
    const int blocks = min((n + threads - 1) / threads, 128);

    int* d_hist;
    cudaMalloc(&d_hist, blocks * RADIX * sizeof(int));

    const float* src = in;
    float* dst = out;

    for (int pass = 0; pass < 4; pass++) {
        int shift = pass * RADIX_BITS;

        cudaMemsetAsync(d_hist, 0, blocks * RADIX * sizeof(int), stream);
        count_radix<<<blocks, threads, 0, stream>>>(src, d_hist, n, shift);
        scan_radix<<<1, RADIX, 0, stream>>>(d_hist, blocks);
        reorder_radix<<<blocks, threads, 0, stream>>>(src, dst, d_hist, n, shift);

        if (pass < 3) {
            src = dst;
            dst = (dst == out) ? tmp : out;
        }
    }

    if (dst != out) {
        cudaMemcpyAsync(out, dst, n * sizeof(float), cudaMemcpyDeviceToDevice, stream);
    }

    cudaFree(d_hist);
}

torch::Tensor radix_sort(torch::Tensor input, torch::Tensor output) {
    auto tmp = torch::empty_like(input);
    radix_sort_impl(
        input.data_ptr<float>(),
        output.data_ptr<float>(),
        tmp.data_ptr<float>(),
        input.numel(),
        c10::cuda::getCurrentCUDAStream()
    );
    return output;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
    m.def("sort", &radix_sort);
}
"""

try:
    radix = load_inline(
        name='radix_sort',
        cpp_sources='',
        cuda_sources=cuda_src,
        functions=['sort'],
        extra_cuda_cflags=['-O3', '--use_fast_math', '-std=c++17'],
        with_cuda=True
    )
    HAS_RADIX = True
except:
    HAS_RADIX = False

def custom_kernel(data):
    input_tensor, output_tensor = data

    if HAS_RADIX:
        radix.sort(input_tensor, output_tensor)
    else:
        torch.sort(input_tensor, out=(output_tensor, torch.empty_like(input_tensor, dtype=torch.long)))

    return output_tensor