submission 71033

  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
-     torch.sort(input_tensor, out=(output_tensor, torch.empty_like(input_tensor, dtype=torch.long)))
+  
+     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