submission 837945

"""impl8 — cheap diagonal detection + fast fill + cuSOLVER dense.

Leaderboard: gpumode `eigh` (LB 775), target GPU B200.

Same structure as the diagonal-fast-path variants, but the diagonal *detection*
is made an order of magnitude cheaper. A matrix is exactly diagonal iff all of
its nonzeros lie on the diagonal, i.e. count_nonzero(A) == count_nonzero(diag(A)).
That's a single pass over A plus a trivial diagonal count -- no `diag_embed`
allocation and no full `A - diag` subtraction (which cost ~4-5 extra passes over
the whole batch). On the n=4096 diagonal benchmark this is the dominant cost, so
shrinking it pulls the (log-sensitive) geomean down noticeably.

Dense solve stays on default cuSOLVER (`torch.linalg.eigh`).
"""
import torch

from task import input_t, output_t


def _is_exact_diagonal(data: torch.Tensor) -> bool:
    diag = torch.diagonal(data, dim1=-2, dim2=-1)
    return torch.count_nonzero(data).item() == torch.count_nonzero(diag).item()


def _diagonal_eigh(data: torch.Tensor) -> output_t:
    batch, n, _ = data.shape
    values, perm = torch.diagonal(data, dim1=-2, dim2=-1).sort(dim=-1)
    vectors = torch.zeros((batch, n, n), device=data.device, dtype=data.dtype)
    vectors.scatter_(1, perm.unsqueeze(1), 1.0)
    return vectors, values.contiguous()


def custom_kernel(data: input_t) -> output_t:
    if data.shape[-1] > 1 and _is_exact_diagonal(data):
        return _diagonal_eigh(data)

    values, vectors = torch.linalg.eigh(data)
    return vectors, values