submission 875138

"""Optimized solution for the `eigh` leaderboard.

Calls cuSOLVER's batched symmetric eigensolvers directly through a small
inline C++ extension instead of going through torch.linalg.eigh:

  * cusolverDnXsyevBatched   -- batched QR-based solver, one call for the
                                whole batch (CUDA >= 12.6u2 / cuSOLVER 11.7.1)
  * cusolverDnSsyevjBatched  -- batched Jacobi solver, kept as an alternative
                                strategy with tunable tolerance / sweeps

Layout trick: a row-major (batch, n, n) float32 buffer reinterpreted as
column-major is A^T, and A is symmetric, so the input buffer can be handed to
cuSOLVER as-is (no transpose copy). On output cuSOLVER stores eigenvectors in
the *columns* of the column-major buffer, which the row-major torch view sees
as the *rows* -- so returning `buf.mT` (a zero-cost view) gives Q in the
required convention.

Robustness: the per-matrix `info` output is checked; if anything failed to
converge the whole batch falls back to torch.linalg.eigh. If the extension
fails to build at all we also fall back, so this file never does worse than
the naive solution correctness-wise.

See EIGH_NOTES.md for the full write-up.
"""

import os

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

CPP_SRC = r"""
#include <torch/extension.h>

#include <c10/cuda/CUDAGuard.h>
#include <cusolverDn.h>

#include <vector>

#if CUSOLVER_VER_MAJOR > 11 ||                                         \
    (CUSOLVER_VER_MAJOR == 11 &&                                       \
     (CUSOLVER_VER_MINOR > 7 ||                                        \
      (CUSOLVER_VER_MINOR == 7 && CUSOLVER_VER_PATCH >= 1)))
#define EIGH_HAVE_XSYEV_BATCHED 1
#else
#define EIGH_HAVE_XSYEV_BATCHED 0
#endif

#define CUSOLVER_CHECK(expr)                                           \
  do {                                                                 \
    cusolverStatus_t status__ = (expr);                                \
    TORCH_CHECK(status__ == CUSOLVER_STATUS_SUCCESS,                   \
                "cuSOLVER error ", static_cast<int>(status__),         \
                " in " #expr);                                         \
  } while (0)

namespace {

cusolverDnHandle_t handle() {
  static cusolverDnHandle_t h = [] {
    cusolverDnHandle_t tmp;
    CUSOLVER_CHECK(cusolverDnCreate(&tmp));
    return tmp;
  }();
  return h;
}

cusolverDnParams_t xt_params() {
  static cusolverDnParams_t p = [] {
    cusolverDnParams_t tmp;
    CUSOLVER_CHECK(cusolverDnCreateParams(&tmp));
    return tmp;
  }();
  return p;
}

std::vector<char>& host_workspace(size_t bytes) {
  static std::vector<char> ws;
  if (ws.size() < bytes) {
    ws.resize(bytes);
  }
  return ws;
}

void check_inputs(const torch::Tensor& A,
                  const torch::Tensor& W,
                  const torch::Tensor& info) {
  TORCH_CHECK(A.is_cuda() && A.scalar_type() == torch::kFloat);
  TORCH_CHECK(A.dim() == 3 && A.size(1) == A.size(2) && A.is_contiguous());
  TORCH_CHECK(W.is_cuda() && W.scalar_type() == torch::kFloat);
  TORCH_CHECK(W.is_contiguous() && W.numel() == A.size(0) * A.size(1));
  TORCH_CHECK(info.is_cuda() && info.scalar_type() == torch::kInt);
  TORCH_CHECK(info.is_contiguous() && info.numel() == A.size(0));
}

}  // namespace

bool has_xsyev_batched() {
  return EIGH_HAVE_XSYEV_BATCHED != 0;
}

// QR-based batched eigensolver. A is overwritten with the eigenvectors
// (column-major per matrix), W receives ascending eigenvalues.
void syev_batched(torch::Tensor A, torch::Tensor W, torch::Tensor info) {
#if EIGH_HAVE_XSYEV_BATCHED
  check_inputs(A, W, info);
  const c10::cuda::CUDAGuard guard(A.device());
  const int64_t batch = A.size(0);
  const int64_t n = A.size(1);

  // No queue redirection: the solver runs on the device's default execution
  // queue, same place the harness runs. The Python caller synchronizes around
  // this call so ordering against torch ops is always safe.
  cusolverDnHandle_t h = handle();

  size_t dev_bytes = 0;
  size_t host_bytes = 0;
  CUSOLVER_CHECK(cusolverDnXsyevBatched_bufferSize(
      h, xt_params(), CUSOLVER_EIG_MODE_VECTOR, CUBLAS_FILL_MODE_LOWER, 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_ws = torch::empty(
      {static_cast<int64_t>(std::max<size_t>(dev_bytes, 16))},
      A.options().dtype(torch::kByte));
  auto& host_ws = host_workspace(host_bytes);

  CUSOLVER_CHECK(cusolverDnXsyevBatched(
      h, xt_params(), CUSOLVER_EIG_MODE_VECTOR, CUBLAS_FILL_MODE_LOWER, n,
      CUDA_R_32F, A.data_ptr<float>(), n,
      CUDA_R_32F, W.data_ptr<float>(), CUDA_R_32F,
      dev_ws.data_ptr(), dev_bytes,
      host_bytes ? host_ws.data() : nullptr, host_bytes,
      info.data_ptr<int>(), batch));
#else
  TORCH_CHECK(false, "cusolverDnXsyevBatched not available in this cuSOLVER");
#endif
}

// Jacobi batched eigensolver, alternative strategy. Same in/out contract as
// syev_batched. Eigenvalue ordering is NOT trusted here; the Python side
// sorts defensively.
void syevj_batched(torch::Tensor A, torch::Tensor W, torch::Tensor info,
                   double tol, int64_t max_sweeps) {
  check_inputs(A, W, info);
  const c10::cuda::CUDAGuard guard(A.device());
  const int batch = static_cast<int>(A.size(0));
  const int n = static_cast<int>(A.size(1));

  cusolverDnHandle_t h = handle();

  syevjInfo_t jp = nullptr;
  CUSOLVER_CHECK(cusolverDnCreateSyevjInfo(&jp));
  CUSOLVER_CHECK(cusolverDnXsyevjSetTolerance(jp, tol));
  CUSOLVER_CHECK(cusolverDnXsyevjSetMaxSweeps(jp, static_cast<int>(max_sweeps)));
  CUSOLVER_CHECK(cusolverDnXsyevjSetSortEig(jp, 1));

  int lwork = 0;
  CUSOLVER_CHECK(cusolverDnSsyevjBatched_bufferSize(
      h, CUSOLVER_EIG_MODE_VECTOR, CUBLAS_FILL_MODE_LOWER, n,
      A.data_ptr<float>(), n, W.data_ptr<float>(), &lwork, jp, batch));

  auto work = torch::empty({std::max(lwork, 1)}, A.options());
  CUSOLVER_CHECK(cusolverDnSsyevjBatched(
      h, CUSOLVER_EIG_MODE_VECTOR, CUBLAS_FILL_MODE_LOWER, n,
      A.data_ptr<float>(), n, W.data_ptr<float>(),
      work.data_ptr<float>(), lwork, info.data_ptr<int>(), jp, batch));

  CUSOLVER_CHECK(cusolverDnDestroySyevjInfo(jp));
}
"""

_EXT = None
_EXT_ERR = None


def _build_ext():
    global _EXT, _EXT_ERR
    if _EXT is not None or _EXT_ERR is not None:
        return _EXT
    try:
        _EXT = load_inline(
            name="eigh_cusolver_batched",
            cpp_sources=CPP_SRC,
            functions=["syev_batched", "syevj_batched", "has_xsyev_batched"],
            with_cuda=True,
            extra_ldflags=["-lcusolver"],
            verbose=False,
        )
    except Exception as err:  # noqa: BLE001 - any build failure => naive fallback
        _EXT_ERR = err
    return _EXT


# Build at import time so compilation never lands inside a timed region.
_build_ext()

# "auto" picks per shape; "syev" / "syevj" / "torch" force one path.
_STRATEGY = os.environ.get("EIGH_STRATEGY", "auto")


def _cusolver_eigh(ext, A: torch.Tensor, strategy: str):
    batch, n = A.shape[0], A.shape[-1]
    # cuSOLVER overwrites its input, and the checker still needs the original
    # A, so one clone is mandatory. Row-major seen column-major = A^T = A.
    buf = A.clone(memory_format=torch.contiguous_format)
    W = torch.empty((batch, n), device=A.device, dtype=torch.float32)
    info = torch.zeros((batch,), device=A.device, dtype=torch.int32)

    # The extension runs on the default CUDA execution queue; synchronize on
    # both sides so it is ordered correctly against surrounding torch ops.
    torch.cuda.synchronize(A.device)
    if strategy == "syevj":
        ext.syevj_batched(buf, W, info, 1e-7, 100)
    else:
        ext.syev_batched(buf, W, info)
    torch.cuda.synchronize(A.device)

    if bool((info != 0).any()):
        # Rare non-convergence: redo the whole batch with the reference path.
        values, vectors = torch.linalg.eigh(A)
        return vectors, values

    if strategy == "syevj":
        # The batched Jacobi solver's sort flag is not trusted blindly:
        # enforce ascending order ourselves (cheap vs the O(n^3) solve).
        W, idx = torch.sort(W, dim=-1)
        buf = torch.take_along_dim(buf, idx.unsqueeze(-1), dim=1)

    # Rows of the row-major view are the eigenvectors => transpose view.
    return buf.mT, W


def custom_kernel(data: input_t) -> output_t:
    A = data
    squeeze = False
    if A.dim() == 2:
        A = A.unsqueeze(0)
        squeeze = True
    batch, n = A.shape[0], A.shape[-1]

    ext = _EXT
    strategy = _STRATEGY
    usable = (
        ext is not None
        and A.is_cuda
        and A.dtype == torch.float32
        and n > 0
    )
    if strategy == "auto":
        if usable and ext.has_xsyev_batched():
            strategy = "syev"
        else:
            strategy = "torch"
    elif not usable:
        strategy = "torch"

    if strategy == "torch":
        values, vectors = torch.linalg.eigh(A)
        Q, L = vectors, values
    else:
        Q, L = _cusolver_eigh(ext, A, strategy)

    if squeeze:
        Q, L = Q.squeeze(0), L.squeeze(0)
    return Q, L