submission 848766

"""Size-dispatched batched symmetric eigendecomposition.

No hacks, default execution context only. Two vendor routines, dispatched by matrix size:

- SMALL n (<= _JACOBI_MAX_N): cusolverDnSsyevjBatched -- cuSOLVER's batched Jacobi solver,
  designed for exactly this regime (many small matrices, one kernel, per-matrix parallelism).
  The ranked score is a GEOMEAN over 13 cases, so the small cases (n=32: 221us, n=176:
  2.96ms, n=352: 17.7ms under the D&C path) count exactly as much as the n=512 monsters --
  and small-n is where the D&C path (tridiagonalization + QR, deep per-matrix pipelines,
  fixed overheads) is furthest from optimal. An earlier syevjBatched test wrote it off as
  "5-8x slower" -- but that was measured at n=512, precisely where Jacobi's O(n^3 * sweeps)
  flop disadvantage bites; the small-n regime it was designed for was never measured.
  Jacobi's tolerance knob also matters: the contest gate is ~200*n*eps (~1e-2 relative at
  n=512) -- orders of magnitude looser than default solver tolerances -- so the tolerance
  is set to 1e-5 (still ~100x tighter than the gate, but allows earlier sweep exit than
  the 1e-7-ish default machine-precision target).

- LARGE n: cusolverDnXsyevBatched (D&C), with the cached-handle + cached-workspace-size
  pattern that took the ranked score from 54.07 -> 50.01ms.

Embedded .cu compiled at import (outside any timed region); falls back to torch.linalg.eigh
on any failure, so worst case is the baseline.
"""

import os
import sys
import tempfile

import torch

input_t = torch.Tensor
output_t = "tuple[torch.Tensor, torch.Tensor]"

_JACOBI_MAX_N = 32    # dispatch threshold: syevjBatched for n <= this, Xsyev D&C above.
                      # Measured on B200 (sub 848724, thresholds at 352): syevj WINS at n=32
                      # (181us vs 221us xsyev / 189us torch) but LOSES at n=176 (11.8ms vs
                      # 2.96ms) and n=352 (29.4ms vs 17.2ms) -- Jacobi's O(n^3*sweeps) flop
                      # disadvantage overtakes its per-matrix parallelism advantage well
                      # before n=176. Keep it only where it wins.
_JACOBI_TOL = 1e-5    # gate is ~1e-2 relative; 1e-5 leaves ~1000x margin yet exits sweeps early
_JACOBI_SWEEPS = 20   # cap; Jacobi on random symmetric converges in ~8 sweeps at these sizes

_CU = r'''
#include <torch/extension.h>
#include <cuda_runtime.h>
#include <cusolverDn.h>
#include <library_types.h>
#include <unordered_map>
#include <stdexcept>
#include <string>
#include <cstdlib>

#define CK(cmd) do { cusolverStatus_t s=cmd; if(s!=CUSOLVER_STATUS_SUCCESS) \
    throw std::runtime_error("cusolver " #cmd " code="+std::to_string(s)); } while(0)

static cusolverDnHandle_t g_handle = nullptr;
static cusolverDnParams_t g_params = nullptr;
static syevjInfo_t g_jparams = nullptr;

struct WsKey { int64_t batch, n, kind; bool operator==(const WsKey& o) const {
    return batch==o.batch && n==o.n && kind==o.kind; } };
struct WsKeyHash { size_t operator()(const WsKey& k) const {
    return (std::hash<int64_t>()(k.batch)*1000003u) ^ std::hash<int64_t>()(k.n) ^ (size_t)k.kind; } };
struct WsEntry { size_t dW, hW; int lwork; };
static std::unordered_map<WsKey, WsEntry, WsKeyHash> g_ws_cache;

static void ensure_handles(double jtol, int jsweeps) {
    if (!g_handle) CK(cusolverDnCreate(&g_handle));
    if (!g_params) CK(cusolverDnCreateParams(&g_params));
    if (!g_jparams) {
        CK(cusolverDnCreateSyevjInfo(&g_jparams));
        CK(cusolverDnXsyevjSetTolerance(g_jparams, jtol));
        CK(cusolverDnXsyevjSetMaxSweeps(g_jparams, jsweeps));
        // sortEig defaults to 1 (ascending), which is exactly the contest convention
    }
}

std::tuple<at::Tensor,at::Tensor> xsyev_batched(at::Tensor A, at::Tensor W) {
    TORCH_CHECK(A.is_cuda() && A.dtype()==at::kFloat && A.dim()==3 && A.size(1)==A.size(2));
    int64_t batch=A.size(0), n=A.size(1);
    ensure_handles(1e-5, 20);
    cusolverEigMode_t jobz=CUSOLVER_EIG_MODE_VECTOR;
    cublasFillMode_t uplo=CUBLAS_FILL_MODE_LOWER;
    WsKey key{batch, n, 0};
    auto it = g_ws_cache.find(key);
    size_t dW, hW;
    if (it == g_ws_cache.end()) {
        dW = 0; hW = 0;
        CK(cusolverDnXsyevBatched_bufferSize(g_handle, g_params, jobz, uplo, n,
            CUDA_R_32F, A.data_ptr<float>(), n, CUDA_R_32F, W.data_ptr<float>(),
            CUDA_R_32F, &dW, &hW, batch));
        g_ws_cache[key] = WsEntry{dW, hW, 0};
    } else {
        dW = it->second.dW; hW = it->second.hW;
    }
    auto dbuf = at::empty({(long)dW}, A.options().dtype(at::kByte));
    void* hbuf = hW ? std::malloc(hW) : nullptr;
    auto info = at::zeros({(long)batch}, A.options().dtype(at::kInt));
    CK(cusolverDnXsyevBatched(g_handle, g_params, jobz, uplo, n,
        CUDA_R_32F, A.data_ptr<float>(), n, CUDA_R_32F, W.data_ptr<float>(),
        CUDA_R_32F, dbuf.data_ptr(), dW, hbuf, hW, info.data_ptr<int>(), batch));
    if (hbuf) std::free(hbuf);
    return std::make_tuple(A, W);
}

std::tuple<at::Tensor,at::Tensor> syevj_batched(at::Tensor A, at::Tensor W, double tol, int64_t sweeps) {
    TORCH_CHECK(A.is_cuda() && A.dtype()==at::kFloat && A.dim()==3 && A.size(1)==A.size(2));
    int64_t batch=A.size(0), n=A.size(1);
    ensure_handles(tol, (int)sweeps);
    cusolverEigMode_t jobz=CUSOLVER_EIG_MODE_VECTOR;
    cublasFillMode_t uplo=CUBLAS_FILL_MODE_LOWER;
    WsKey key{batch, n, 1};
    auto it = g_ws_cache.find(key);
    int lwork;
    if (it == g_ws_cache.end()) {
        lwork = 0;
        CK(cusolverDnSsyevjBatched_bufferSize(g_handle, jobz, uplo, (int)n,
            A.data_ptr<float>(), (int)n, W.data_ptr<float>(), &lwork, g_jparams, (int)batch));
        g_ws_cache[key] = WsEntry{0, 0, lwork};
    } else {
        lwork = it->second.lwork;
    }
    auto work = at::empty({(long)lwork}, A.options());
    auto info = at::zeros({(long)batch}, A.options().dtype(at::kInt));
    CK(cusolverDnSsyevjBatched(g_handle, jobz, uplo, (int)n,
        A.data_ptr<float>(), (int)n, W.data_ptr<float>(),
        work.data_ptr<float>(), lwork, info.data_ptr<int>(), g_jparams, (int)batch));
    return std::make_tuple(A, W);
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m){
    m.def("xsyev_batched", &xsyev_batched);
    m.def("syevj_batched", &syevj_batched);
}
'''

_ext = None
_failed = False


def _load():
    global _ext, _failed
    if _ext is not None or _failed:
        return _ext
    try:
        import torch.utils.cpp_extension as cpp
        d = tempfile.mkdtemp(prefix="eigh_dispatch_")
        src = os.path.join(d, "x.cu")
        with open(src, "w") as f:
            f.write(_CU)
        _ext = cpp.load(name="eigh_dispatch_ext", sources=[src],
                        extra_ldflags=["-lcusolver"], verbose=False)
    except Exception as e:
        print(f"[dispatch] compile failed -> torch.eigh: {e}", file=sys.stderr)
        _failed = True
        _ext = None
    return _ext


if torch.cuda.is_available():
    _load()


def kernel(data: input_t):
    A = data
    if not (torch.is_tensor(A) and A.is_cuda and A.dtype == torch.float32
            and A.dim() == 3 and A.shape[-1] == A.shape[-2]):
        res = torch.linalg.eigh(A)
        return res.eigenvectors, res.eigenvalues
    ext = _load()
    if ext is None:
        res = torch.linalg.eigh(A.contiguous())
        return res.eigenvectors.to(torch.float32), res.eigenvalues.to(torch.float32)
    B, n, _ = A.shape
    Aw = A.contiguous().clone()            # overwritten with eigenvectors (col-major)
    W = torch.empty(B, n, dtype=torch.float32, device=A.device)
    try:
        if n <= _JACOBI_MAX_N:
            ext.syevj_batched(Aw, W, _JACOBI_TOL, _JACOBI_SWEEPS)
            # finiteness guard only on the sweep-capped Jacobi path (the sync it forces is
            # a few us -- noise at ms scale, but keep it off the big D&C cases anyway)
            if not bool(torch.isfinite(W).all()):
                raise RuntimeError("non-finite eigenvalues")
        else:
            ext.xsyev_batched(Aw, W)
    except Exception as e:
        print(f"[dispatch] runtime fail -> torch.eigh: {e}", file=sys.stderr)
        res = torch.linalg.eigh(A.contiguous())
        return res.eigenvectors.to(torch.float32), res.eigenvalues.to(torch.float32)
    Q = Aw.transpose(-1, -2).contiguous()  # col-major V read row-major = V^T
    return Q, W


custom_kernel = kernel
solve = kernel