submission 150629

  ⋯ 1 unchanged lines
   
  from task import input_t, output_t
  import torch
- import triton
- import triton.language as tl
+ from torch.utils.cpp_extension import load_inline
   
+ # CUB-based radix sort implementation
+ cuda_source = """
+ #include <torch/extension.h>
+ #include <cub/device/device_radix_sort.cuh>
+ #include <cuda_runtime.h>
+ #include <c10/cuda/CUDACachingAllocator.h>
   
- @triton.jit
- def simple_sort_kernel(
-     input_ptr,
-     output_ptr,
-     n_elements,
-     BLOCK_SIZE: tl.constexpr,
- ):
-     """Simple sort kernel using tl.sort."""
-     offsets = tl.arange(0, BLOCK_SIZE)
-     mask = offsets < n_elements
-     data = tl.load(input_ptr + offsets, mask=mask, other=float('inf'))
+ void cub_radix_sort_float(
+     const float* keys_in,
+     float* keys_out,
+     int64_t n) {
   
-     # Use Triton's built-in sort
-     sorted_data = tl.sort(data)
+     // Determine temporary device storage requirements
+     void* d_temp_storage = nullptr;
+     size_t temp_storage_bytes = 0;
   
-     tl.store(output_ptr + offsets, sorted_data, mask=mask)
+     cub::DeviceRadixSort::SortKeys(
+         d_temp_storage, temp_storage_bytes,
+         keys_in, keys_out, n);
   
+     // Allocate temporary storage using PyTorch's caching allocator
+     auto& allocator = *c10::cuda::CUDACachingAllocator::get();
+     auto temp_storage = allocator.allocate(temp_storage_bytes);
+     d_temp_storage = temp_storage.get();
   
- def triton_custom_kernel(data: input_t) -> output_t:
-     inp, output = data
+     // Run sorting operation
+     cub::DeviceRadixSort::SortKeys(
+         d_temp_storage, temp_storage_bytes,
+         keys_in, keys_out, n);
+     // temp_storage is automatically freed when it goes out of scope
+ }
   
-     n_elements = inp.numel()
+ torch::Tensor sort_kernel_cuda(torch::Tensor input, torch::Tensor output) {
+     TORCH_CHECK(input.is_cuda(), "input must be a CUDA tensor");
+     TORCH_CHECK(output.is_cuda(), "output must be a CUDA tensor");
+     TORCH_CHECK(input.is_contiguous(), "input must be contiguous");
+     TORCH_CHECK(output.is_contiguous(), "output must be contiguous");
+     TORCH_CHECK(input.dtype() == torch::kFloat32, "input must be float32");
   
-     # Round up to next power of 2 for sort
-     block_size = 1
-     while block_size < n_elements:
-         block_size *= 2
+     int64_t n = input.numel();
   
-     # Cap block size to reasonable maximum
-     block_size = min(block_size, 1024)
+     const float* input_ptr = input.data_ptr<float>();
+     float* output_ptr = output.data_ptr<float>();
   
-     # Launch kernel
-     simple_sort_kernel[(1,)](
-         inp,
-         output,
-         n_elements,
-         BLOCK_SIZE=block_size,
-     )
+     cub_radix_sort_float(input_ptr, output_ptr, n);
   
+     return output;
+ }
+ """
+  
+ cpp_source = """
+ torch::Tensor sort_kernel_cuda(torch::Tensor input, torch::Tensor output);
+ """
+  
+ # Compile the CUDA extension
+ cub_sort_module = load_inline(
+     name='cub_sort',
+     cpp_sources=cpp_source,
+     cuda_sources=cuda_source,
+     functions=['sort_kernel_cuda'],
+     with_cuda=True,
+     extra_cuda_cflags=['-O3', '-use_fast_math'],
+ )
+  
+ def cub_kernel(data: input_t) -> output_t:
+     inp, output = data
+     cub_sort_module.sort_kernel_cuda(inp, output)
      return output
   
  def naive_custom_kernel(data):
  ⋯ 3 unchanged lines
      return out
   
   
- custom_kernel = naive_custom_kernel
+ custom_kernel = cub_kernel