submission 851285

"""Batched eigh via cuSOLVER's cusolverDnXsyevBatched (CUDA >= 12.6).

This is the batched, 64-bit-API version of syevd: LAPACK-quality QR/DC
reduction, one call for the entire batch — the kernel PyTorch's eigh loop
never uses (it calls syevd per matrix). Bound straight out of libcusolver
with ctypes; no compilation, nothing outside the standard environment.

Column-major trick: a row-major symmetric matrix reinterpreted column-major
is itself, so the input transfers as-is; the eigenvector output comes back
transposed, so Q = out.transpose(-1, -2). Eigenvalues are ascending, matching
torch.linalg.eigh.

Fallbacks: torch.linalg.eigh for CPU tensors, n <= 32 (torch already batches
those), missing symbol (older CUDA), nonzero status/info. Input never mutated.
"""

import ctypes
import glob
import os

import torch
from task import input_t, output_t

_DIRECT_N = 32
_CUSOLVER_EIG_MODE_VECTOR = 1
_CUBLAS_FILL_MODE_LOWER = 0
_CUDA_R_32F = 0


def _find_cusolver() -> ctypes.CDLL | None:
    candidates: list[str] = []
    try:
        import nvidia.cusolver

        candidates += glob.glob(
            os.path.join(os.path.dirname(nvidia.cusolver.__file__), "lib", "libcusolver.so*")
        )
    except Exception:
        pass
    torch_dir = os.path.dirname(torch.__file__)
    candidates += glob.glob(os.path.join(torch_dir, "lib", "libcusolver*"))
    candidates += glob.glob(
        os.path.join(os.path.dirname(torch_dir), "nvidia", "cusolver", "lib", "libcusolver.so*")
    )
    candidates += ["libcusolver.so.12", "libcusolver.so.13", "libcusolver.so.11", "libcusolver.so"]
    for path in candidates:
        try:
            return ctypes.CDLL(path)
        except OSError:
            continue
    return None


class _XsyevBatched:
    def __init__(self) -> None:
        lib = _find_cusolver()
        if lib is None:
            raise OSError("libcusolver not found")
        self.lib = lib
        p, i32, i64, sz = ctypes.c_void_p, ctypes.c_int, ctypes.c_int64, ctypes.c_size_t
        lib.cusolverDnCreate.argtypes = [ctypes.POINTER(p)]
        lib.cusolverDnCreateParams.argtypes = [ctypes.POINTER(p)]
        lib.cusolverDnXsyevBatched_bufferSize.argtypes = [
            p, p, i32, i32, i64, i32, p, i64, i32, p, i32,
            ctypes.POINTER(sz), ctypes.POINTER(sz), i64,
        ]
        lib.cusolverDnXsyevBatched.argtypes = [
            p, p, i32, i32, i64, i32, p, i64, i32, p, i32,
            p, sz, p, sz, p, i64,
        ]

        self.handle = p()
        self._check(lib.cusolverDnCreate(ctypes.byref(self.handle)), "create")
        self.params = p()
        self._check(lib.cusolverDnCreateParams(ctypes.byref(self.params)), "params")
        self._dev_work: dict[tuple[int, int], torch.Tensor] = {}
        self._host_work: dict[tuple[int, int], ctypes.Array] = {}

    @staticmethod
    def _check(status: int, what: str) -> None:
        if status != 0:
            raise RuntimeError(f"cusolver {what} failed: status {status}")

    def __call__(self, a: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """a: (batch, n, n) fp32 contiguous CUDA tensor, overwritten with
        eigenvectors (column-major). Returns (transposed vectors, values)."""
        lib = self.lib
        batch, n, _ = a.shape
        w = torch.empty((batch, n), device=a.device, dtype=torch.float32)
        info = torch.zeros(batch, device=a.device, dtype=torch.int32)

        key = (n, batch)
        dev_work = self._dev_work.get(key)
        host_work = self._host_work.get(key)
        if dev_work is None:
            dev_bytes = ctypes.c_size_t(0)
            host_bytes = ctypes.c_size_t(0)
            self._check(
                lib.cusolverDnXsyevBatched_bufferSize(
                    self.handle, self.params,
                    _CUSOLVER_EIG_MODE_VECTOR, _CUBLAS_FILL_MODE_LOWER,
                    n, _CUDA_R_32F, ctypes.c_void_p(a.data_ptr()), n,
                    _CUDA_R_32F, ctypes.c_void_p(w.data_ptr()), _CUDA_R_32F,
                    ctypes.byref(dev_bytes), ctypes.byref(host_bytes), batch,
                ),
                "bufferSize",
            )
            dev_work = torch.empty(
                max(int(dev_bytes.value), 16), device=a.device, dtype=torch.uint8
            )
            host_work = ctypes.create_string_buffer(max(int(host_bytes.value), 16))
            self._dev_work[key] = dev_work
            self._host_work[key] = host_work

        self._check(
            lib.cusolverDnXsyevBatched(
                self.handle, self.params,
                _CUSOLVER_EIG_MODE_VECTOR, _CUBLAS_FILL_MODE_LOWER,
                n, _CUDA_R_32F, ctypes.c_void_p(a.data_ptr()), n,
                _CUDA_R_32F, ctypes.c_void_p(w.data_ptr()), _CUDA_R_32F,
                ctypes.c_void_p(dev_work.data_ptr()), dev_work.numel(),
                ctypes.cast(host_work, ctypes.c_void_p), len(host_work),
                ctypes.c_void_p(info.data_ptr()), batch,
            ),
            "XsyevBatched",
        )
        if bool((info != 0).any()):
            raise RuntimeError("XsyevBatched: nonzero info")
        return a.transpose(-1, -2), w


_solver: _XsyevBatched | None = None
_solver_failed = False


def _get_solver() -> _XsyevBatched | None:
    global _solver, _solver_failed
    if _solver is None and not _solver_failed:
        try:
            _solver = _XsyevBatched()
        except Exception:
            _solver_failed = True
    return _solver


def custom_kernel(data: input_t) -> output_t:
    batch, n, _ = data.shape

    solver = _get_solver() if data.is_cuda and n > _DIRECT_N else None
    if solver is None:
        values, vectors = torch.linalg.eigh(data)
        return vectors, values

    scale = data.abs().amax(dim=(-2, -1), keepdim=True)
    scale = scale.clamp_min(torch.finfo(torch.float32).tiny)
    aw = (data / scale).contiguous()

    try:
        q, w = solver(aw)
    except Exception:
        values, vectors = torch.linalg.eigh(data)
        return vectors, values

    return q, w * scale.squeeze(-1)