submission 862867

#!POPCORN leaderboard eigh
#!POPCORN gpu B200

import ctypes
import glob
from pathlib import Path

import torch
from task import input_t, output_t


_CUSOLVER = None
_HANDLE = ctypes.c_void_p()
_SYEVJ_PARAMS = ctypes.c_void_p()
_LWORK_CACHE = {}

CUSOLVER_EIG_MODE_VECTOR = 1
CUBLAS_FILL_MODE_UPPER = 1


def _check(status: int) -> None:
    if status != 0:
        raise RuntimeError(f"cuSOLVER error: {status}")


def _init_cusolver():
    global _CUSOLVER
    if _CUSOLVER is not None:
        return _CUSOLVER

    candidates = [
        "libcusolver.so",
        "libcusolver.so.13",
        "libcusolver.so.12",
    ]
    torch_lib = Path(torch.__file__).resolve().parent / "lib"
    candidates.extend(glob.glob("/usr/local/cuda*/lib64/libcusolver.so*"))
    candidates.extend(glob.glob(str(torch_lib / "libcusolver.so*")))

    last_error = None
    for candidate in candidates:
        try:
            lib = ctypes.CDLL(candidate)
            break
        except OSError as exc:
            last_error = exc
    else:
        raise last_error or OSError("could not load libcusolver")

    lib.cusolverDnCreate.argtypes = [ctypes.POINTER(ctypes.c_void_p)]
    lib.cusolverDnCreate.restype = ctypes.c_int
    lib.cusolverDnDestroy.argtypes = [ctypes.c_void_p]
    lib.cusolverDnDestroy.restype = ctypes.c_int

    lib.cusolverDnCreateSyevjInfo.argtypes = [ctypes.POINTER(ctypes.c_void_p)]
    lib.cusolverDnCreateSyevjInfo.restype = ctypes.c_int
    lib.cusolverDnDestroySyevjInfo.argtypes = [ctypes.c_void_p]
    lib.cusolverDnDestroySyevjInfo.restype = ctypes.c_int

    lib.cusolverDnSsyevjBatched_bufferSize.argtypes = [
        ctypes.c_void_p,
        ctypes.c_int,
        ctypes.c_int,
        ctypes.c_int,
        ctypes.c_void_p,
        ctypes.c_int,
        ctypes.c_void_p,
        ctypes.POINTER(ctypes.c_int),
        ctypes.c_void_p,
        ctypes.c_int,
    ]
    lib.cusolverDnSsyevjBatched_bufferSize.restype = ctypes.c_int

    lib.cusolverDnSsyevjBatched.argtypes = [
        ctypes.c_void_p,
        ctypes.c_int,
        ctypes.c_int,
        ctypes.c_int,
        ctypes.c_void_p,
        ctypes.c_int,
        ctypes.c_void_p,
        ctypes.c_void_p,
        ctypes.c_int,
        ctypes.c_void_p,
        ctypes.c_void_p,
        ctypes.c_int,
    ]
    lib.cusolverDnSsyevjBatched.restype = ctypes.c_int

    _check(lib.cusolverDnCreate(ctypes.byref(_HANDLE)))
    _check(lib.cusolverDnCreateSyevjInfo(ctypes.byref(_SYEVJ_PARAMS)))
    _CUSOLVER = lib
    return lib


def _syevj_batched(data: torch.Tensor) -> output_t:
    lib = _init_cusolver()
    batch = data.shape[0]
    n = data.shape[-1]
    a = data.clone()
    values = torch.empty((batch, n), device=data.device, dtype=torch.float32)
    info = torch.empty((batch,), device=data.device, dtype=torch.int32)

    cache_key = (batch, n)
    lwork = _LWORK_CACHE.get(cache_key)
    if lwork is None:
        lwork_c = ctypes.c_int()
        _check(
            lib.cusolverDnSsyevjBatched_bufferSize(
                _HANDLE,
                CUSOLVER_EIG_MODE_VECTOR,
                CUBLAS_FILL_MODE_UPPER,
                n,
                ctypes.c_void_p(a.data_ptr()),
                n,
                ctypes.c_void_p(values.data_ptr()),
                ctypes.byref(lwork_c),
                _SYEVJ_PARAMS,
                batch,
            )
        )
        lwork = int(lwork_c.value)
        _LWORK_CACHE[cache_key] = lwork

    work = torch.empty((lwork,), device=data.device, dtype=torch.float32)
    _check(
        lib.cusolverDnSsyevjBatched(
            _HANDLE,
            CUSOLVER_EIG_MODE_VECTOR,
            CUBLAS_FILL_MODE_UPPER,
            n,
            ctypes.c_void_p(a.data_ptr()),
            n,
            ctypes.c_void_p(values.data_ptr()),
            ctypes.c_void_p(work.data_ptr()),
            lwork,
            ctypes.c_void_p(info.data_ptr()),
            _SYEVJ_PARAMS,
            batch,
        )
    )
    return a.transpose(-1, -2), values


def custom_kernel(data: input_t) -> output_t:
    if data.shape[0] == 20 and data.shape[-1] == 32:
        return _syevj_batched(data)
    values, vectors = torch.linalg.eigh(data)
    return vectors, values