submission 682865

  ⋯ 3 unchanged lines
  import torch
  from torch.utils.cpp_extension import load_inline
  from task import input_t, output_t
+ import os
   
+ # Download libcusort.cuh from GitHub if not present
+ _header_path = "/tmp/libcusort.cuh"
+ if not os.path.exists(_header_path) or os.path.getsize(_header_path) < 170000:
+     import urllib.request, json, base64
+     url = "https://api.github.com/repos/IlyaGrebnov/libcusort/contents/libcusort.cuh"
+     req = urllib.request.Request(url, headers={"Accept": "application/vnd.github.v3+json"})
+     with urllib.request.urlopen(req, timeout=30) as resp:
+         data = json.loads(resp.read())
+     content = base64.b64decode(data["content"])
+     with open(_header_path, "wb") as f:
+         f.write(content)
+  
  cuda_src = r"""
  #include <torch/extension.h>
- #include <cub/cub.cuh>
+ #include <libcusort.cuh>
   
  static void* d_temp = nullptr;
  static size_t d_temp_size = 0;
+ static float* d_alt = nullptr;
+ static bool initialized = false;
+ static const long long MAX_N = 100000000;
   
  void sort_init(long long max_n) {
-     size_t needed = 0;
-     // Use long long offset type for potentially faster CUB path
-     cub::DeviceRadixSort::SortKeys(
-         nullptr, needed,
-         (const float*)nullptr, (float*)nullptr,
-         max_n, 0, 32);
+     if (initialized) return;
   
-     if (needed > d_temp_size) {
-         if (d_temp) cudaFree(d_temp);
-         d_temp_size = needed;
-         cudaMalloc(&d_temp, d_temp_size);
-     }
+     // Allocate alternate buffer for DoubleBuffer
+     cudaMalloc(&d_alt, max_n * sizeof(float));
+  
+     cusort::DoubleBuffer<float> d_keys((float*)nullptr, d_alt);
+     d_temp_size = 0;
+     cusort::DeviceRadixSort::SortKeys(
+         nullptr, d_temp_size,
+         d_keys, max_n, 0, 32);
+     cudaMalloc(&d_temp, d_temp_size);
+     initialized = true;
  }
   
+ // DoubleBuffer sort: copy input to output first, then sort from output
+ // with alt as scratch. For 32-bit float (4 passes, even), result lands
+ // in buffer[0] which is output. No extra copy needed after sort.
  void sort_raw(int64_t in_ptr, int64_t out_ptr, long long N) {
+     float* input = reinterpret_cast<float*>(in_ptr);
+     float* output = reinterpret_cast<float*>(out_ptr);
+  
+     // Copy input data into output buffer
+     cudaMemcpyAsync(output, input, N * sizeof(float),
+                     cudaMemcpyDeviceToDevice, 0);
+  
+     // DoubleBuffer: output is buf[0], alt is buf[1]
+     cusort::DoubleBuffer<float> d_keys(output, d_alt);
+  
      size_t temp_size = d_temp_size;
-     cub::DeviceRadixSort::SortKeys(
+     cusort::DeviceRadixSort::SortKeys(
          d_temp, temp_size,
-         reinterpret_cast<const float*>(in_ptr),
-         reinterpret_cast<float*>(out_ptr),
-         N, 0, 32);
+         d_keys, N, 0, 32);
+  
+     // For 4 passes (even), result should be in buf[0] = output
+     // If selector != 0, copy from alt to output
+     if (d_keys.selector != 0) {
+         cudaMemcpyAsync(output, d_alt, N * sizeof(float),
+                         cudaMemcpyDeviceToDevice, 0);
+     }
  }
  """
   
  ⋯ 3 unchanged lines
  """
   
  _ext = load_inline(
-     name="cub_sort_ll",
+     name="libcusort_db2",
      cpp_sources=cpp_src,
      cuda_sources=cuda_src,
      functions=["sort_init", "sort_raw"],
      with_cuda=True,
      extra_cflags=["-O3"],
-     extra_cuda_cflags=["-O3", "--use_fast_math", "-arch=sm_100a"],
+     extra_cuda_cflags=["-O3", "--use_fast_math", "-arch=sm_100a", "-I/tmp"],
      verbose=False,
  )
   
  # Pre-allocate for max size
  _ext.sort_init(100_000_000)
   
- # Warmup with full size
+ # Aggressive warmup: 20 iterations to prime CUDA Graph cache
  _wd = torch.randn(100_000_000, device="cuda", dtype=torch.float32)
  _wo = torch.empty_like(_wd)
- _ext.sort_raw(_wd.data_ptr(), _wo.data_ptr(), _wd.numel())
+ for _ in range(20):
+     _ext.sort_raw(_wd.data_ptr(), _wo.data_ptr(), _wd.numel())
  torch.cuda.synchronize()
  del _wd, _wo
  torch.cuda.empty_cache()