submission 571333

import hashlib
import os
import shutil
import subprocess
import threading
from pathlib import Path
from typing import Dict, Optional

import torch

from task import input_t, output_t


_CPP_TEMPLATE = r'''
#include <torch/library.h>
#include <ATen/ATen.h>
#include <cstdint>

int64_t cub_sort_temp_storage_bytes_cuda(const at::Tensor& input);
at::Tensor cub_sort_out_cuda(const at::Tensor& input, const at::Tensor& output, const at::Tensor& scratch);

TORCH_LIBRARY(__NAMESPACE__, m) {
  m.def("temp_storage_bytes(Tensor input) -> int");
  m.def("sort_out(Tensor input, Tensor output, Tensor scratch) -> Tensor");
}

TORCH_LIBRARY_IMPL(__NAMESPACE__, CUDA, m) {
  m.impl("temp_storage_bytes", cub_sort_temp_storage_bytes_cuda);
  m.impl("sort_out", cub_sort_out_cuda);
}
'''


_CU_TEMPLATE = r'''
#include <ATen/ATen.h>
#include <c10/cuda/CUDAGuard.h>
#include <c10/cuda/CUDAException.h>

#include <cub/device/device_radix_sort.cuh>

#include <cstdint>
#include <limits>

namespace {

void check_tensor_args(const at::Tensor& input, const at::Tensor& output, const at::Tensor& scratch) {
  TORCH_CHECK(input.is_cuda(), "input must be CUDA");
  TORCH_CHECK(output.is_cuda(), "output must be CUDA");
  TORCH_CHECK(scratch.is_cuda(), "scratch must be CUDA");
  TORCH_CHECK(input.scalar_type() == at::kFloat, "input must be float32");
  TORCH_CHECK(output.scalar_type() == at::kFloat, "output must be float32");
  TORCH_CHECK(scratch.scalar_type() == at::kByte, "scratch must be uint8");
  TORCH_CHECK(input.dim() == 1, "input must be 1D");
  TORCH_CHECK(output.dim() == 1, "output must be 1D");
  TORCH_CHECK(input.is_contiguous(), "input must be contiguous");
  TORCH_CHECK(output.is_contiguous(), "output must be contiguous");
  TORCH_CHECK(scratch.is_contiguous(), "scratch must be contiguous");
  TORCH_CHECK(input.numel() == output.numel(), "input/output size mismatch");
  TORCH_CHECK(input.get_device() == output.get_device(), "input/output device mismatch");
  TORCH_CHECK(input.get_device() == scratch.get_device(), "scratch must be on same device as input");
  TORCH_CHECK(input.data_ptr<float>() != output.data_ptr<float>(), "input and output must be distinct tensors");
}

size_t required_temp_storage_bytes(int64_t num_items) {
  TORCH_CHECK(num_items >= 0, "num_items must be non-negative");
  TORCH_CHECK(num_items <= std::numeric_limits<int>::max(), "num_items exceeds CUB int limit");
  size_t temp_storage_bytes = 0;
  C10_CUDA_CHECK(cub::DeviceRadixSort::SortKeys(
      nullptr,
      temp_storage_bytes,
      static_cast<const float*>(nullptr),
      static_cast<float*>(nullptr),
      static_cast<int>(num_items)));
  return temp_storage_bytes;
}

}  // namespace

int64_t cub_sort_temp_storage_bytes_cuda(const at::Tensor& input) {
  TORCH_CHECK(input.is_cuda(), "input must be CUDA");
  TORCH_CHECK(input.scalar_type() == at::kFloat, "input must be float32");
  TORCH_CHECK(input.dim() == 1, "input must be 1D");
  return static_cast<int64_t>(required_temp_storage_bytes(input.numel()));
}

at::Tensor cub_sort_out_cuda(const at::Tensor& input, const at::Tensor& output, const at::Tensor& scratch) {
  check_tensor_args(input, output, scratch);

  const c10::cuda::CUDAGuard device_guard(input.device());
  size_t temp_storage_bytes = required_temp_storage_bytes(input.numel());
  const size_t scratch_bytes = static_cast<size_t>(scratch.numel()) * scratch.element_size();
  TORCH_CHECK(scratch_bytes >= temp_storage_bytes,
              "scratch too small: have ", scratch_bytes, " bytes but need ", temp_storage_bytes);

  void* temp_storage = scratch.data_ptr();
  const float* keys_in = input.data_ptr<float>();
  float* keys_out = output.data_ptr<float>();

  C10_CUDA_CHECK(cub::DeviceRadixSort::SortKeys(
      temp_storage,
      temp_storage_bytes,
      keys_in,
      keys_out,
      static_cast<int>(input.numel()),
      0,
      static_cast<int>(sizeof(float) * 8)));

  return output;
}
'''


_BUILD_LOCK = threading.Lock()
_BUILD_HASH: Optional[str] = None
_NAMESPACE: Optional[str] = None
_SORT_OUT = None
_TEMP_STORAGE_BYTES_OP = None
_TEMP_STORAGE_BYTES_CACHE: Dict[int, int] = {}
_SCRATCH_BY_DEVICE: Dict[int, torch.Tensor] = {}


def _cuda_include_dir(cuda_home: Path) -> Path:
    candidates = [
        cuda_home / "targets" / "x86_64-linux" / "include",
        cuda_home / "include",
    ]
    for path in candidates:
        if path.exists():
            return path
    raise RuntimeError(f"Could not find CUDA include directory under {cuda_home}")


def _cuda_lib_dir(cuda_home: Path) -> Path:
    candidates = [
        cuda_home / "targets" / "x86_64-linux" / "lib",
        cuda_home / "lib64",
        cuda_home / "lib",
    ]
    for path in candidates:
        if path.exists():
            return path
    raise RuntimeError(f"Could not find CUDA library directory under {cuda_home}")


def _run_checked(command: list[str], cwd: Path) -> None:
    proc = subprocess.run(command, cwd=str(cwd), capture_output=True, text=True)
    if proc.returncode != 0:
        joined = " ".join(command)
        stdout = proc.stdout[-12000:]
        stderr = proc.stderr[-12000:]
        raise RuntimeError(
            f"Command failed: {joined}\nstdout:\n{stdout}\nstderr:\n{stderr}"
        )


def _build_sources(namespace: str) -> tuple[str, str]:
    return (
        _CPP_TEMPLATE.replace("__NAMESPACE__", namespace),
        _CU_TEMPLATE.replace("__NAMESPACE__", namespace),
    )


def _build_metadata() -> tuple[str, str, Path, str, int, tuple[int, int]]:
    from torch.utils.cpp_extension import include_paths, library_paths

    cuda_home = Path(os.environ.get("CUDA_HOME", "/usr/local/cuda"))
    torch_includes = include_paths()
    torch_libs = library_paths()
    if not torch_includes or not torch_libs:
        raise RuntimeError("Failed to discover PyTorch include/library paths")

    abi = int(torch._C._GLIBCXX_USE_CXX11_ABI)
    capability = torch.cuda.get_device_capability()
    hash_material = "\n".join(
        [
            _CPP_TEMPLATE,
            _CU_TEMPLATE,
            torch.__version__,
            str(torch.version.cuda),
            str(abi),
            str(capability),
            str(torch_includes),
            str(torch_libs),
            str(cuda_home),
        ]
    )
    build_hash = hashlib.sha256(hash_material.encode("utf-8")).hexdigest()[:16]
    namespace = f"glia_sort_{build_hash}"
    build_dir = Path("/tmp") / namespace
    return namespace, build_hash, build_dir, str(cuda_home), abi, capability


def _build_shared_library() -> Path:
    namespace, build_hash, build_dir, cuda_home_str, abi, capability = _build_metadata()

    cpp_source, cu_source = _build_sources(namespace)
    build_dir.mkdir(parents=True, exist_ok=True)

    cpp_path = build_dir / f"{namespace}.cpp"
    cu_path = build_dir / f"{namespace}.cu"
    cpp_obj = build_dir / f"{namespace}.o"
    cu_obj = build_dir / f"{namespace}_cuda.o"
    lib_path = build_dir / f"{namespace}.so"
    tmp_lib_path = build_dir / f"{namespace}.so.tmp"

    if lib_path.exists():
        return lib_path

    cpp_path.write_text(cpp_source)
    cu_path.write_text(cu_source)

    from torch.utils.cpp_extension import include_paths, library_paths

    torch_include_paths = include_paths()
    torch_lib_dir = Path(library_paths()[0])
    cuda_home = Path(cuda_home_str)
    cuda_include_dir = _cuda_include_dir(cuda_home)
    cuda_lib_dir = _cuda_lib_dir(cuda_home)

    gpp = shutil.which("g++") or "/usr/bin/g++"
    nvcc = str((cuda_home / "bin" / "nvcc").resolve())

    includes = [f"-I{path}" for path in torch_include_paths]
    includes.append(f"-I{cuda_include_dir}")
    abi_flag = f"-D_GLIBCXX_USE_CXX11_ABI={abi}"
    arch = f"{capability[0]}{capability[1]}"


    _run_checked(
        [
            gpp,
            "-c",
            "-O3",
            "-fPIC",
            "-std=c++17",
            abi_flag,
            *includes,
            str(cpp_path),
            "-o",
            str(cpp_obj),
        ],
        build_dir,
    )

    _run_checked(
        [
            nvcc,
            "-c",
            "-O3",
            "--std=c++17",
            "-Xcompiler",
            "-fPIC",
            abi_flag,
            *includes,
            f"-gencode=arch=compute_{arch},code=sm_{arch}",
            str(cu_path),
            "-o",
            str(cu_obj),
        ],
        build_dir,
    )

    _run_checked(
        [
            gpp,
            "-shared",
            "-O3",
            "-fPIC",
            str(cpp_obj),
            str(cu_obj),
            "-o",
            str(tmp_lib_path),
            f"-L{torch_lib_dir}",
            "-ltorch",
            "-ltorch_cpu",
            "-lc10",
            "-ltorch_cuda",
            "-lc10_cuda",
            f"-L{cuda_lib_dir}",
            "-lcudart",
            f"-Wl,-rpath,{torch_lib_dir}",
            f"-Wl,-rpath,{cuda_lib_dir}",
        ],
        build_dir,
    )

    os.replace(tmp_lib_path, lib_path)
    return lib_path


def _ensure_loaded() -> None:
    global _BUILD_HASH, _NAMESPACE, _SORT_OUT, _TEMP_STORAGE_BYTES_OP
    if _SORT_OUT is not None:
        return

    with _BUILD_LOCK:
        if _SORT_OUT is not None:
            return

        namespace, build_hash, _, _, _, _ = _build_metadata()
        lib_path = _build_shared_library()
        torch.ops.load_library(str(lib_path))
        namespace_obj = getattr(torch.ops, namespace)

        _BUILD_HASH = build_hash
        _NAMESPACE = namespace
        _SORT_OUT = namespace_obj.sort_out
        _TEMP_STORAGE_BYTES_OP = namespace_obj.temp_storage_bytes


def _required_temp_storage_bytes(input_tensor: torch.Tensor) -> int:
    num_items = int(input_tensor.numel())
    cached = _TEMP_STORAGE_BYTES_CACHE.get(num_items)
    if cached is not None:
        return cached
    required = int(_TEMP_STORAGE_BYTES_OP(input_tensor))
    _TEMP_STORAGE_BYTES_CACHE[num_items] = required
    return required


def _scratch_buffer(device: torch.device, required_bytes: int) -> torch.Tensor:
    device_index = 0 if device.index is None else int(device.index)
    scratch = _SCRATCH_BY_DEVICE.get(device_index)
    if scratch is None or int(scratch.numel()) < required_bytes:
        scratch = torch.empty(required_bytes, device=device, dtype=torch.uint8)
        _SCRATCH_BY_DEVICE[device_index] = scratch
    return scratch


def custom_kernel(data: input_t) -> output_t:
    input_tensor, output_tensor = data
    _ensure_loaded()
    required_bytes = _required_temp_storage_bytes(input_tensor)
    scratch = _scratch_buffer(input_tensor.device, required_bytes)
    return _SORT_OUT(input_tensor, output_tensor, scratch)