submission 877467

#!POPCORN leaderboard eigh
#!POPCORN gpu B200

# Per-shape routing between cuSOLVER batched solvers:
#  - small n: batched Jacobi (syevjBatched), tolerance scaled to the n-dependent
#    correctness gate (gate = 200*n*eps, so small n needs a tight Jacobi tol)
#  - large n: batched divide-and-conquer (cusolverDnXsyevBatched)

import torch
from torch.utils.cpp_extension import load_inline
from task import input_t, output_t

CPP_SRC = r"""
#include <torch/extension.h>
torch::Tensor syevj_batched(torch::Tensor A, torch::Tensor W, double tol, int64_t max_sweeps);
torch::Tensor syev_batched(torch::Tensor A, torch::Tensor W);
"""

CUDA_SRC = r"""
#include <torch/extension.h>
#include <cusolverDn.h>
#include <vector>

#define CUSOLVER_CHECK(expr) do { \
    cusolverStatus_t st_ = (expr); \
    TORCH_CHECK(st_ == CUSOLVER_STATUS_SUCCESS, "cusolver error ", (int)st_); \
} while (0)

static cusolverDnHandle_t get_handle() {
    static cusolverDnHandle_t handle = nullptr;
    if (!handle) CUSOLVER_CHECK(cusolverDnCreate(&handle));
    return handle;
}

// Batched Jacobi. A is overwritten with eigenvectors (column-major convention).
torch::Tensor syevj_batched(torch::Tensor A, torch::Tensor W, double tol, int64_t max_sweeps) {
    const int batch = A.size(0);
    const int n = A.size(1);
    cusolverDnHandle_t handle = get_handle();

    syevjInfo_t params;
    CUSOLVER_CHECK(cusolverDnCreateSyevjInfo(&params));
    CUSOLVER_CHECK(cusolverDnXsyevjSetTolerance(params, tol));
    CUSOLVER_CHECK(cusolverDnXsyevjSetMaxSweeps(params, (int)max_sweeps));
    CUSOLVER_CHECK(cusolverDnXsyevjSetSortEig(params, 1));

    auto info = torch::empty({batch}, A.options().dtype(torch::kInt32));
    int lwork = 0;
    CUSOLVER_CHECK(cusolverDnSsyevjBatched_bufferSize(
        handle, CUSOLVER_EIG_MODE_VECTOR, CUBLAS_FILL_MODE_UPPER, n,
        A.data_ptr<float>(), n, W.data_ptr<float>(), &lwork, params, batch));
    auto work = torch::empty({lwork}, A.options());
    CUSOLVER_CHECK(cusolverDnSsyevjBatched(
        handle, CUSOLVER_EIG_MODE_VECTOR, CUBLAS_FILL_MODE_UPPER, n,
        A.data_ptr<float>(), n, W.data_ptr<float>(), work.data_ptr<float>(), lwork,
        info.data_ptr<int>(), params, batch));
    CUSOLVER_CHECK(cusolverDnDestroySyevjInfo(params));
    return info;
}

// Batched divide-and-conquer (generic 64-bit API). A overwritten with
// eigenvectors (column-major convention), W ascending eigenvalues.
torch::Tensor syev_batched(torch::Tensor A, torch::Tensor W) {
    const int64_t batch = A.size(0);
    const int64_t n = A.size(1);
    cusolverDnHandle_t handle = get_handle();

    static cusolverDnParams_t params = nullptr;
    if (!params) CUSOLVER_CHECK(cusolverDnCreateParams(&params));

    auto info = torch::empty({batch}, A.options().dtype(torch::kInt32));
    size_t dev_bytes = 0, host_bytes = 0;
    CUSOLVER_CHECK(cusolverDnXsyevBatched_bufferSize(
        handle, params, CUSOLVER_EIG_MODE_VECTOR, CUBLAS_FILL_MODE_UPPER, n,
        CUDA_R_32F, A.data_ptr<float>(), n,
        CUDA_R_32F, W.data_ptr<float>(),
        CUDA_R_32F, &dev_bytes, &host_bytes, batch));
    auto dev_work = torch::empty({(int64_t)dev_bytes}, A.options().dtype(torch::kUInt8));
    std::vector<uint8_t> host_work(host_bytes ? host_bytes : 1);
    CUSOLVER_CHECK(cusolverDnXsyevBatched(
        handle, params, CUSOLVER_EIG_MODE_VECTOR, CUBLAS_FILL_MODE_UPPER, n,
        CUDA_R_32F, A.data_ptr<float>(), n,
        CUDA_R_32F, W.data_ptr<float>(),
        CUDA_R_32F, dev_work.data_ptr(), dev_bytes,
        host_work.data(), host_bytes,
        info.data_ptr<int>(), batch));
    return info;
}
"""

_mod = load_inline(
    name="eigh_final",
    cpp_sources=[CPP_SRC],
    cuda_sources=[CUDA_SRC],
    functions=["syevj_batched", "syev_batched"],
    extra_ldflags=["-lcusolver"],
    verbose=True,
)

# Routing (from per-case B200 benchmarks, see PROGRESS.md):
#  - n <= JACOBI_MAX_N: batched Jacobi with tolerance tight enough for the
#    n-scaled gate (gate = 200*n*eps; e.g. 7.6e-4 at n=32)
#  - n >= DC_MIN_N: our batched divide-and-conquer binding (ties/beats torch)
#  - otherwise: torch.linalg.eigh (best measured on mid shapes)
JACOBI_MAX_N = 32
JACOBI_TOL_TIGHT = 1e-5
JACOBI_SWEEPS = 100
DC_MIN_N = 1 << 30  # D&C path kept but disabled: ties torch within runner noise


# --- fast path for +/-1-clustered spectra -------------------------------
# The `clustered` generator glues all eigenvalues to -1 or +1 with ~1e-5
# jitter, far inside the n-scaled gates. Then (I+A)/2 spans the +1 invariant
# subspace, so both eigenspaces can be extracted with randomized range
# finding + CholeskyQR (all batched GEMMs) instead of a full eigensolve.
# Detection is a cheap probe of ||A(Av) - v||, required to hold for EVERY
# matrix in the batch (mixed batches fall back to the general solver).

_PROBE_REL_TOL = 1e-3
_gen = torch.Generator(device="cuda")
_gen.manual_seed(0x5EED)


def _chol_qr(y: torch.Tensor, ridge: float) -> torch.Tensor:
    g = y.transpose(-1, -2) @ y
    if ridge > 0.0:
        scale = torch.diagonal(g, dim1=-2, dim2=-1).mean(-1, keepdim=True)
        g = g + torch.diag_embed(ridge * scale.expand(-1, g.shape[-1]))
    ell = torch.linalg.cholesky(g)
    return torch.linalg.solve_triangular(
        ell.transpose(-1, -2), y, upper=True, left=False
    )


def _complete_basis(u: torch.Tensor) -> torch.Tensor:
    # Given orthonormal u (batch, n, r), return orthonormal (batch, n, n)
    # whose first r columns span span(u), via Householder reconstruction:
    # with D chosen so M = I - U_top @ D is diagonally biased and M = P L R,
    # H = I - V T V^T is exactly orthogonal with H[:, :r] = u @ D.
    batch, n, r = u.shape
    dev, dt = u.device, u.dtype
    u_top = u[:, :r, :]
    u_bot = u[:, r:, :]
    d = torch.diagonal(u_top, dim1=-2, dim2=-1).sign()
    d = torch.where(d == 0, torch.ones_like(d), d).neg().unsqueeze(-2)

    m = -(u_top * d)
    m.diagonal(dim1=-2, dim2=-1).add_(1.0)
    lu, pivots = torch.linalg.lu_factor(m)
    perm, ell, r_up = torch.lu_unpack(lu, pivots)

    v_top = perm @ ell
    v_bot = torch.linalg.solve_triangular(r_up, -(u_bot * d), upper=True, left=False)
    v = torch.cat([v_top, v_bot], dim=-2)
    # T = R (P L)^{-T}  solved as  T L^T = R P  =>  T = (R P) L^{-T}
    t = torch.linalg.solve_triangular(
        ell.transpose(-1, -2), r_up @ perm, upper=True, left=False
    )

    w = t @ v_bot.transpose(-1, -2)          # (batch, r, n-r)
    comp = -(v @ w)                          # (batch, n, n-r)
    comp.diagonal(offset=-r, dim1=-2, dim2=-1).add_(1.0)
    return torch.cat([u * d, comp], dim=-1)


def _try_pm1_clustered(a: torch.Tensor):
    batch, n, _ = a.shape
    v = torch.randn((n, 2), device=a.device, dtype=a.dtype, generator=_gen)
    z = a @ (a @ v)
    rel = (z - v).abs().amax(dim=(-2, -1)) / v.abs().max()
    if not bool((rel < _PROBE_REL_TOL).all().item()):
        return None

    # eigenvalues are +/-1: multiplicity of +1 from the trace, batch-uniform.
    tr = (n + torch.diagonal(a, dim1=-2, dim2=-1).sum(-1)) * 0.5
    r_plus = int(tr.mean().round().item())
    if not bool(((tr - r_plus).abs() < 0.25).all().item()):
        return None
    r_minus = n - r_plus

    if r_minus == 0 or r_plus == 0:
        # A ~ +/-I: the identity basis is an exact eigenbasis.
        lam = torch.diagonal(a, dim1=-2, dim2=-1)
        lam, order = torch.sort(lam, dim=-1)
        eye = torch.eye(n, device=a.device, dtype=a.dtype).expand(batch, n, n)
        q = torch.gather(eye, -1, order.unsqueeze(-2).expand(batch, n, n))
        return q.contiguous(), lam.contiguous()

    # Range-find only the SMALLER eigenspace, then complete the basis exactly
    # (the orthogonal complement IS the other eigenspace). The projected
    # Gaussian is effectively square in the target subspace, so a single
    # CholeskyQR can rotate the span on ill-conditioned draws — subspace
    # iteration (re-applying the exact projector) purges that error and
    # leaves a well-conditioned basis for the final clean passes.
    r = min(r_minus, r_plus)
    sign = -1.0 if r_minus <= r_plus else 1.0
    omega = torch.randn((n, r), device=a.device, dtype=a.dtype, generator=_gen)
    y = omega + sign * (a @ omega)          # 2 * P @ omega
    q1 = _chol_qr(y, ridge=1e-3)
    y = q1 + sign * (a @ q1)                # re-project: purge span error
    q1 = _chol_qr(y, ridge=1e-6)
    q1 = _chol_qr(q1, ridge=0.0)
    q = _complete_basis(q1)

    # self-verify against the actual gates using the A@Q we need anyway;
    # any surprise routes to the safe solver (perf gamble, never correctness)
    aq = a @ q
    lam = (q * aq).sum(dim=-2)
    resid = (aq - q * lam.unsqueeze(-2)).abs().sum(dim=-2).amax(dim=-1)
    a_norm = a.abs().sum(dim=-2).amax(dim=-1)
    gate = (200.0 * n * torch.finfo(torch.float32).eps) * a_norm
    gq = q.transpose(-1, -2) @ q
    gq.diagonal(dim1=-2, dim2=-1).sub_(1.0)
    orth_err = gq.abs().sum(dim=-2).amax(dim=-1)
    orth_gate = 100.0 * n * torch.finfo(torch.float32).eps
    ok = (resid < 0.5 * gate).all() & (orth_err < 0.5 * orth_gate).all()
    if not bool(ok.item()):
        return None

    lam, order = torch.sort(lam, dim=-1)
    q = torch.gather(q, -1, order.unsqueeze(-2).expand(batch, n, n))
    return q.contiguous(), lam.contiguous()


def custom_kernel(data: input_t) -> output_t:
    n = data.shape[-1]
    if n <= JACOBI_MAX_N or n >= DC_MIN_N:
        a = data.clone()
        w = torch.empty(data.shape[:-1], dtype=data.dtype, device=data.device)
        if n <= JACOBI_MAX_N:
            _mod.syevj_batched(a, w, JACOBI_TOL_TIGHT, JACOBI_SWEEPS)
        else:
            _mod.syev_batched(a, w)
        return a.transpose(-1, -2), w
    if n >= 256:
        try:
            fast = _try_pm1_clustered(data)
        except Exception:
            fast = None
        if fast is not None:
            return fast
    values, vectors = torch.linalg.eigh(data)
    return vectors, values