pytorch_fullgraph_compiled_kernel.py

"""
CompiledEigh: torch.compile(fullgraph=True) drop-in for torch.linalg.eigh.

Eliminates all graph breaks, device-host syncs, and dynamic allocation.
Output contract matches torch.linalg.eigh:
  eigenvalues:  [*, n] real, ascending
  eigenvectors: [*, n, n] orthonormal columns

Author: AbstractPhil / GeoLIP project
"""

import math
import torch
import torch.nn as nn
from torch import Tensor
from typing import Tuple, Optional


# =============================================================================
# Constants
# =============================================================================

DEFAULT_MAX_NEWTON: int = 8
DEFAULT_MAX_JACOBI_SWEEPS: int = 10  # 10 sweeps saturates for n <= 16
JACOBI_THRESHOLD: int = 16


# =============================================================================
# Atom: 2x2 Symmetric Eigenproblem
# =============================================================================

def eigh_2x2(a: Tensor, b: Tensor, c: Tensor, eps: float = 1e-30
             ) -> Tuple[Tensor, Tensor, Tensor, Tensor]:
    """
    Closed-form eigendecomposition of batched 2x2 symmetric matrices.
    Returns: (lambda1, lambda2, cos_theta, sin_theta), lambda1 <= lambda2.
    """
    trace = a + c
    diff = a - c
    two_b = 2.0 * b
    hyp = torch.sqrt(diff * diff + two_b * two_b + eps)

    lambda1 = 0.5 * (trace - hyp)
    lambda2 = 0.5 * (trace + hyp)

    vx = two_b
    vy = lambda2 - a
    norm_v = torch.sqrt(vx * vx + vy * vy + eps)
    cos_theta = vy / norm_v
    sin_theta = vx / norm_v

    return lambda1, lambda2, cos_theta, sin_theta


# =============================================================================
# Utility: Newton-Schulz Orthogonalization (all-bmm, GPU-native)
# =============================================================================

def orthogonalize_ns(V: Tensor, n_iter: int = 2) -> Tensor:
    """
    Re-orthogonalize columns of V via Newton-Schulz iteration.

    Computes V @ (V^T V)^{-1/2} using the coupled iteration:
        X_0 = I,  Y_0 = V^T V
        X_{k+1} = 0.5 * X_k @ (3I - Y_k)
        Y_{k+1} = 0.5 * (3I - Y_k) @ Y_k

    X converges to (V^T V)^{-1/2}, Y converges to I.
    Cubically convergent when V^T V ≈ I.

    Convergence from ||V^T V - I|| = ε:
        1 iteration: error → O(ε²)  ≈ 1e-6 from 1e-3
        2 iterations: error → O(ε⁴) ≈ 1e-12 from 1e-3

    All ops are bmm — fully compiled, no sequential column processing.

    V: [B, n, n] (square, columns are approximate eigenvectors)
    Returns: [B, n, n] with orthonormal columns
    """
    B, n, m = V.shape
    I_n = torch.eye(m, device=V.device, dtype=V.dtype).unsqueeze(0).expand(B, -1, -1)

    # Z = V^T V ≈ I
    Y = torch.bmm(V.transpose(-2, -1), V)
    X = I_n.clone()

    for _ in range(n_iter):
        T = 3.0 * I_n - Y           # (3I - Y_k)
        X = 0.5 * torch.bmm(X, T)   # X_{k+1}
        Y = 0.5 * torch.bmm(T, Y)   # Y_{k+1}

    return torch.bmm(V, X)


# =============================================================================
# Phase 1: Householder Tridiagonalization
# =============================================================================

class HouseholderTridiagonalizer(nn.Module):
    """
    Reduces batched symmetric A to tridiagonal T = Q^T A Q via
    Householder reflections. Fixed loop bounds, compilable.
    """

    def __init__(self, max_n: int, eps: float = 1e-30):
        super().__init__()
        self.max_n = max_n
        self.eps = eps

    def forward(self, A: Tensor, d: Tensor, e: Tensor,
                reflectors: Tensor) -> None:
        B, n, _ = A.shape
        eps = self.eps

        for k in range(n - 2):
            tail_len = n - k - 1
            x = A[:, k + 1:, k].clone()

            sigma = torch.sqrt((x * x).sum(dim=-1, keepdim=True) + eps)
            sign_x0 = torch.where(x[:, 0:1] >= 0,
                                  torch.ones_like(sigma),
                                  -torch.ones_like(sigma))
            alpha = -sign_x0 * sigma

            v = x.clone()
            v[:, 0:1] = v[:, 0:1] - alpha
            v_norm = torch.sqrt((v * v).sum(dim=-1, keepdim=True) + eps)
            v = v / v_norm

            reflectors[k, :, :tail_len] = v
            if tail_len < n:
                reflectors[k, :, tail_len:] = 0.0

            sub_A = A[:, k + 1:, k + 1:]
            v_col = v.unsqueeze(-1)

            p = torch.bmm(sub_A, v_col).squeeze(-1)
            vtp = (v * p).sum(dim=-1, keepdim=True)
            q = p - vtp * v

            q_col = q.unsqueeze(-1)
            q_row = q.unsqueeze(-2)
            v_row = v.unsqueeze(-2)

            A[:, k + 1:, k + 1:] -= 2.0 * (v_col @ q_row + q_col @ v_row)
            A[:, k, k + 1] = alpha.squeeze(-1)
            A[:, k + 1, k] = alpha.squeeze(-1)

        for i in range(n):
            d[:, i] = A[:, i, i]
        for i in range(n - 1):
            e[:, i] = A[:, i, i + 1]


# =============================================================================
# Phase 2a: Secular Equation Newton Solver (Fixed Budget)
# =============================================================================

class SecularNewtonSolver(nn.Module):

    def __init__(self, max_newton: int = DEFAULT_MAX_NEWTON,
                 eps: float = 1e-30, tol: float = 1e-7):
        super().__init__()
        self.max_newton = max_newton
        self.eps = eps
        self.tol = tol

    def forward(self, delta: Tensor, z_sq: Tensor,
                rho: Tensor, mask: Tensor) -> Tensor:
        B, m = delta.shape
        eps = self.eps
        tol = self.tol

        z_sq_sum = (z_sq * mask).sum(dim=-1, keepdim=True)
        rho_abs = rho.abs().unsqueeze(-1)
        upper_bound = delta[:, -1:] + z_sq_sum * rho_abs + 1.0

        lo = delta + eps
        hi = torch.cat([delta[:, 1:], upper_bound], dim=-1) - eps
        lam = 0.5 * (lo + hi)
        rho_exp = rho.unsqueeze(-1)

        for _step in range(self.max_newton):
            delta_exp = delta.unsqueeze(-1)
            lam_exp = lam.unsqueeze(-2)
            denom = delta_exp - lam_exp

            denom_safe = torch.where(
                denom.abs() < eps,
                torch.full_like(denom, eps) * denom.sign().clamp(min=0.5),
                denom
            )

            z_sq_exp = z_sq.unsqueeze(-1)
            mask_exp = mask.unsqueeze(-1)
            masked_z = z_sq_exp * mask_exp

            terms = masked_z / denom_safe
            f = 1.0 + rho_exp * terms.sum(dim=-2)
            f_prime = rho_exp * (masked_z / (denom_safe * denom_safe)).sum(dim=-2)

            f_prime_safe = torch.where(
                f_prime.abs() < eps,
                torch.full_like(f_prime, eps),
                f_prime
            )
            delta_lam = -f / f_prime_safe
            lam_new = torch.clamp(lam + delta_lam, lo, hi)

            f_pos = f > 0
            lo = torch.where(f_pos & mask.bool(), lam, lo)
            hi = torch.where(~f_pos & mask.bool(), lam, hi)

            converged = (f.abs() < tol) | ~mask.bool()
            lam = torch.where(converged, lam, lam_new)

        return lam


# =============================================================================
# Phase 2b: Eigenvectors from Secular Equation
# =============================================================================

def secular_eigenvectors(delta: Tensor, lam: Tensor, z: Tensor,
                         mask: Tensor, eps: float = 1e-30) -> Tensor:
    delta_exp = delta.unsqueeze(-1)
    lam_exp = lam.unsqueeze(-2)
    denom = delta_exp - lam_exp

    denom_safe = torch.where(
        denom.abs() < eps,
        torch.full_like(denom, eps) * denom.sign().clamp(min=0.5),
        denom
    )

    z_exp = z.unsqueeze(-1)
    mask_exp = mask.unsqueeze(-1)
    V = (z_exp * mask_exp) / denom_safe
    col_norms = torch.sqrt((V * V).sum(dim=-2, keepdim=True) + eps)
    V = V / col_norms
    return V


# =============================================================================
# Phase 2c: Fixed-Depth Tensor Tree D&C
# =============================================================================

class TensorTreeDC(nn.Module):

    def __init__(self, max_n: int,
                 max_newton: int = DEFAULT_MAX_NEWTON,
                 eps: float = 1e-30, tol: float = 1e-7):
        super().__init__()
        self.padded_n = 1 << math.ceil(math.log2(max(max_n, 2)))
        self.depth = int(math.log2(self.padded_n))
        self.max_n = max_n
        self.eps = eps
        self.secular_solver = SecularNewtonSolver(
            max_newton=max_newton, eps=eps, tol=tol
        )

    def forward(self, d: Tensor, e: Tensor) -> Tuple[Tensor, Tensor]:
        B, n = d.shape
        pn = self.padded_n
        eps = self.eps
        device = d.device
        dtype = d.dtype

        if n < pn:
            d_max = d.abs().max(dim=-1, keepdim=True).values + 1.0
            pad_diag = d_max + torch.arange(1, pn - n + 1, device=device, dtype=dtype).unsqueeze(0)
            d_padded = torch.cat([d, pad_diag], dim=-1)
            e_padded = torch.zeros(B, pn - 1, device=device, dtype=dtype)
            e_padded[:, :n - 1] = e
        else:
            d_padded = d.clone()
            e_padded = e.clone()

        # DOWNWARD PASS
        coupling_rho = []
        current_d = d_padded.clone()
        current_e = e_padded.clone()

        for level in range(self.depth):
            num_sub = 2 ** level
            sub_size = pn // num_sub
            half = sub_size // 2

            cd = current_d.reshape(B, num_sub, sub_size)
            ce = current_e.reshape(B, num_sub, sub_size - 1)

            rho = ce[:, :, half - 1].clone()
            coupling_rho.append(rho)

            cd[:, :, half - 1] = cd[:, :, half - 1] - rho.abs()
            cd[:, :, half] = cd[:, :, half] - rho.abs()
            ce[:, :, half - 1] = 0.0

            left_d = cd[:, :, :half].reshape(B, num_sub * half)
            right_d = cd[:, :, half:].reshape(B, num_sub * half)
            current_d = torch.stack([left_d.reshape(B, num_sub, half),
                                     right_d.reshape(B, num_sub, half)],
                                    dim=2).reshape(B, pn)

            left_e = ce[:, :, :half - 1].reshape(B, num_sub, half - 1)
            right_e = ce[:, :, half:].reshape(B, num_sub, half - 1)
            current_e = torch.stack([left_e, right_e], dim=2).reshape(
                B, num_sub * 2 * (half - 1))
            expected_e_len = pn - 1
            if current_e.shape[-1] < expected_e_len:
                current_e = torch.nn.functional.pad(
                    current_e, (0, expected_e_len - current_e.shape[-1]))

        # BASE
        base_evals = current_d
        V_current = torch.ones(B, pn, 1, 1, device=device, dtype=dtype)

        # UPWARD PASS
        current_evals = base_evals

        for level in range(self.depth - 1, -1, -1):
            num_sub = 2 ** level
            sub_size = pn // num_sub
            half = sub_size // 2
            child_size = half

            evals_grouped = current_evals.reshape(B, num_sub, 2, child_size)
            left_evals = evals_grouped[:, :, 0, :]
            right_evals = evals_grouped[:, :, 1, :]
            delta = torch.cat([left_evals, right_evals], dim=-1)

            V_grouped = V_current.reshape(B, num_sub, 2, child_size, child_size)
            V_left = V_grouped[:, :, 0, :, :]
            V_right = V_grouped[:, :, 1, :, :]

            z_left = V_left[:, :, -1, :]
            z_right = V_right[:, :, 0, :]
            z_cat = torch.cat([z_left, z_right], dim=-1)

            delta_sorted, sort_idx = delta.sort(dim=-1)
            z_sorted = z_cat.gather(-1, sort_idx)

            rho = coupling_rho[level]
            mask = torch.ones(B, num_sub, sub_size, device=device, dtype=dtype)

            gaps = (delta_sorted[:, :, 1:] - delta_sorted[:, :, :-1]).abs()
            degenerate = gaps < (eps * 100)
            avg = 0.5 * (delta_sorted[:, :, :-1] + delta_sorted[:, :, 1:])

            delta_defl = delta_sorted.clone()
            delta_defl[:, :, :-1] = torch.where(degenerate, avg, delta_sorted[:, :, :-1])
            delta_defl[:, :, 1:] = torch.where(degenerate, avg, delta_sorted[:, :, 1:])

            z_defl = z_sorted.clone()
            defl_kill = torch.ones_like(z_sorted)
            defl_kill[:, :, 1:] = torch.where(
                degenerate, torch.zeros_like(gaps), torch.ones_like(gaps))
            z_defl = z_defl * defl_kill

            z_sq = z_defl * z_defl
            Bns = B * num_sub
            new_evals_flat = self.secular_solver(
                delta_defl.reshape(Bns, sub_size),
                z_sq.reshape(Bns, sub_size),
                rho.reshape(Bns),
                mask.reshape(Bns, sub_size),
            )
            new_evals = new_evals_flat.reshape(B, num_sub, sub_size)

            V_secular_flat = secular_eigenvectors(
                delta_defl.reshape(Bns, sub_size),
                new_evals_flat,
                z_defl.reshape(Bns, sub_size),
                mask.reshape(Bns, sub_size),
                eps=eps
            )
            V_secular = V_secular_flat.reshape(B, num_sub, sub_size, sub_size)

            inv_sort = sort_idx.argsort(dim=-1)
            inv_exp = inv_sort.unsqueeze(-1).expand_as(V_secular)
            V_unsorted = V_secular.gather(-2, inv_exp)

            V_block = torch.zeros(B, num_sub, sub_size, sub_size,
                                  device=device, dtype=dtype)
            V_block[:, :, :half, :half] = V_left
            V_block[:, :, half:, half:] = V_right

            V_merged = torch.bmm(
                V_block.reshape(Bns, sub_size, sub_size),
                V_unsorted.reshape(Bns, sub_size, sub_size)
            ).reshape(B, num_sub, sub_size, sub_size)

            current_evals = new_evals.reshape(B, pn)
            V_current = V_merged

        eigenvalues = current_evals
        eigenvectors = V_current.squeeze(1)

        sorted_evals, sort_perm = eigenvalues.sort(dim=-1)
        sort_exp = sort_perm.unsqueeze(-2).expand_as(eigenvectors)
        sorted_evecs = eigenvectors.gather(-1, sort_exp)

        if n < pn:
            sorted_evals = sorted_evals[:, :n]
            sorted_evecs = sorted_evecs[:, :n, :n]

        return sorted_evals, sorted_evecs


# =============================================================================
# Phase 2 (alternate): Jacobi for small n
# =============================================================================

class JacobiEigh(nn.Module):
    """
    Jacobi eigenvalue algorithm for small symmetric matrices.
    Fixed sweep count, fully vectorized, zero branches.

    COMPILE FIX: Pair indices stored as plain Python lists (not tensors).
    Dynamo sees these as constants — no SymInt issues.
    """

    def __init__(self, max_n: int,
                 max_sweeps: int = DEFAULT_MAX_JACOBI_SWEEPS,
                 eps: float = 1e-30):
        super().__init__()
        self.max_n = max_n
        self.max_sweeps = max_sweeps
        self.eps = eps

        # CRITICAL: plain Python lists, NOT registered buffers.
        # Dynamo traces these as compile-time constants.
        pairs = []
        for p in range(max_n):
            for q in range(p + 1, max_n):
                pairs.append((p, q))
        self._pairs_p: list[int] = [p for p, q in pairs]
        self._pairs_q: list[int] = [q for p, q in pairs]
        self._n_pairs: int = len(pairs)

    def forward(self, A: Tensor) -> Tuple[Tensor, Tensor]:
        """
        A: [B, n, n] symmetric
        Returns: (eigenvalues [B, n] ascending, eigenvectors [B, n, n])
        """
        B, n, _ = A.shape
        eps = self.eps

        W = A.clone()
        V = torch.eye(n, device=A.device, dtype=A.dtype).unsqueeze(0).expand(B, -1, -1).clone()

        for _sweep in range(self.max_sweeps):
            for idx in range(self._n_pairs):
                # Plain Python ints — Dynamo sees these as constants
                p: int = self._pairs_p[idx]
                q: int = self._pairs_q[idx]

                app = W[:, p, p]
                aqq = W[:, q, q]
                apq = W[:, p, q]

                # Givens rotation angle
                two_apq = 2.0 * apq
                diff = aqq - app

                # Safe division: sign-preserving eps guard
                abs_two_apq = two_apq.abs().clamp(min=eps)
                sign_two_apq = torch.where(two_apq >= 0,
                                           torch.ones_like(two_apq),
                                           -torch.ones_like(two_apq))
                tau = diff / (abs_two_apq * sign_two_apq)

                tau_sign = torch.where(tau >= 0,
                                       torch.ones_like(tau),
                                       -torch.ones_like(tau))
                t = tau_sign / (tau.abs() + torch.sqrt(1.0 + tau * tau))

                # Zero rotation when off-diagonal is already negligible
                skip = (apq.abs() < eps).float()
                t = t * (1.0 - skip)

                c = 1.0 / torch.sqrt(1.0 + t * t)
                s = t * c

                # ── Rotate W columns p, q ──
                Wp = W[:, :, p].clone()
                Wq = W[:, :, q].clone()
                c_col = c.unsqueeze(-1)
                s_col = s.unsqueeze(-1)
                W[:, :, p] = c_col * Wp - s_col * Wq
                W[:, :, q] = s_col * Wp + c_col * Wq

                # ── Rotate W rows p, q ──
                Wp = W[:, p, :].clone()
                Wq = W[:, q, :].clone()
                W[:, p, :] = c_col * Wp - s_col * Wq
                W[:, q, :] = s_col * Wp + c_col * Wq

                # ── Exact diagonal repair (prevents accumulation drift) ──
                W[:, p, q] = 0.0
                W[:, q, p] = 0.0
                W[:, p, p] = app - t * apq
                W[:, q, q] = aqq + t * apq

                # ── Accumulate eigenvectors ──
                Vp = V[:, :, p].clone()
                Vq = V[:, :, q].clone()
                V[:, :, p] = c_col * Vp - s_col * Vq
                V[:, :, q] = s_col * Vp + c_col * Vq

        # ── Newton-Schulz re-orthogonalization ──
        # 2 iterations: orth error 1e-3 → ~1e-12 via bmm (GPU-native)
        V = orthogonalize_ns(V, n_iter=2)

        # ── Extract and sort ──
        eigenvalues = torch.diagonal(W, dim1=-2, dim2=-1)
        sorted_evals, sort_perm = eigenvalues.sort(dim=-1)
        sort_exp = sort_perm.unsqueeze(-2).expand_as(V)
        sorted_evecs = V.gather(-1, sort_exp)

        return sorted_evals, sorted_evecs


# =============================================================================
# Phase 3: Householder Back-Accumulation
# =============================================================================

class HouseholderBackAccumulate(nn.Module):

    def __init__(self, max_n: int, eps: float = 1e-30):
        super().__init__()
        self.max_n = max_n
        self.eps = eps

    def forward(self, reflectors: Tensor, Z: Tensor, n: int) -> Tensor:
        V = Z.clone()
        for k in range(n - 3, -1, -1):
            tail_len = n - k - 1
            v = reflectors[k, :, :tail_len]
            v_col = v.unsqueeze(-1)
            V_sub = V[:, k + 1:, :]
            vtV = torch.bmm(v_col.transpose(-2, -1), V_sub)
            V[:, k + 1:, :] = V_sub - 2.0 * v_col @ vtV
        return V


# =============================================================================
# Validation
# =============================================================================

class EighValidator(nn.Module):

    def forward(self, A: Tensor, eigenvalues: Tensor,
                eigenvectors: Tensor) -> Tuple[Tensor, Tensor, Tensor]:
        B, n, _ = A.shape
        AV = torch.bmm(A, eigenvectors)
        VL = eigenvectors * eigenvalues.unsqueeze(-2)
        residual = AV - VL
        A_norm = torch.linalg.norm(A.reshape(B, -1), dim=-1).clamp(min=1e-30)
        residual_norm = torch.linalg.norm(residual.reshape(B, -1), dim=-1) / A_norm

        VtV = torch.bmm(eigenvectors.transpose(-2, -1), eigenvectors)
        I = torch.eye(n, device=A.device, dtype=A.dtype).unsqueeze(0)
        orth_err = torch.linalg.norm((VtV - I).reshape(B, -1), dim=-1)

        return residual_norm, orth_err, residual_norm.max()


# =============================================================================
# Top-Level: CompiledEigh
# =============================================================================

class CompiledEigh(nn.Module):
    """
    Drop-in replacement for torch.linalg.eigh.

    Usage:
        solver = CompiledEigh(max_n=6)
        solver = torch.compile(solver, fullgraph=True)
        eigenvalues, eigenvectors = solver(A)
    """

    def __init__(self, max_n: int,
                 use_jacobi: Optional[bool] = None,
                 max_newton: int = DEFAULT_MAX_NEWTON,
                 max_jacobi_sweeps: int = DEFAULT_MAX_JACOBI_SWEEPS,
                 eps: float = 1e-30, tol: float = 1e-7):
        super().__init__()
        self.max_n = max_n
        self.eps = eps

        if use_jacobi is None:
            use_jacobi = (max_n <= JACOBI_THRESHOLD)
        self.use_jacobi = use_jacobi

        if use_jacobi:
            self.jacobi = JacobiEigh(
                max_n=max_n, max_sweeps=max_jacobi_sweeps, eps=eps)
        else:
            self.tridiag = HouseholderTridiagonalizer(max_n=max_n, eps=eps)
            self.dc = TensorTreeDC(
                max_n=max_n, max_newton=max_newton, eps=eps, tol=tol)
            self.back_accum = HouseholderBackAccumulate(max_n=max_n, eps=eps)

        self.validator = EighValidator()

    def forward(self, A: Tensor, validate: bool = False
                ) -> Tuple[Tensor, Tensor]:
        B, n, _ = A.shape

        if self.use_jacobi:
            eigenvalues, eigenvectors = self.jacobi(A)
        else:
            A_work = A.clone()
            d = torch.empty(B, n, device=A.device, dtype=A.dtype)
            e = torch.empty(B, n - 1, device=A.device, dtype=A.dtype)
            reflectors = torch.zeros(max(n - 2, 1), B, n,
                                     device=A.device, dtype=A.dtype)
            self.tridiag(A_work, d, e, reflectors)
            eigenvalues, Z = self.dc(d, e)
            eigenvectors = self.back_accum(reflectors, Z, n)
            # Newton-Schulz re-orthogonalization for D&C path
            eigenvectors = orthogonalize_ns(eigenvectors, n_iter=2)

        if validate:
            res_norm, orth_err, max_err = self.validator(A, eigenvalues, eigenvectors)
            print(f"[CompiledEigh] max residual: {max_err.item():.2e}, "
                  f"mean orth err: {orth_err.mean().item():.2e}")

        return eigenvalues, eigenvectors


# =============================================================================
# Functional API
# =============================================================================

_cached_solvers = {}

def compiled_eigh(A: Tensor, validate: bool = False) -> Tuple[Tensor, Tensor]:
    B, n, _ = A.shape
    key = (n, A.device, A.dtype)
    if key not in _cached_solvers:
        _cached_solvers[key] = CompiledEigh(max_n=n).to(A.device)
    return _cached_solvers[key](A, validate=validate)


"""
CompiledEigh — Colab GPU Benchmark v3
Fixes:
  v2: Jacobi pairs as plain Python lists (Dynamo compile fix), sweeps 6→10
  v3: Replaced Gram-Schmidt with Newton-Schulz orthogonalization (all-bmm),
      disabled TF32 to ensure fp32 precision on Blackwell
"""

import torch
import time
import gc
import sys

# ── Ensure full fp32 precision on Ampere/Hopper/Blackwell ──
# TF32 uses 10-bit mantissa for matmul which can degrade orthogonality
torch.backends.cuda.matmul.allow_tf32 = False
torch.backends.cudnn.allow_tf32 = False
torch.set_float32_matmul_precision('highest')


def sync():
    if torch.cuda.is_available():
        torch.cuda.synchronize()


def gpu_timer(fn, warmup=10, repeats=200):
    for _ in range(warmup):
        fn()
    sync()
    start = time.perf_counter()
    for _ in range(repeats):
        fn()
    sync()
    return (time.perf_counter() - start) / repeats


def make_symmetric_batch(B, n, device, dtype=torch.float32):
    R = torch.randn(B, n, n, device=device, dtype=dtype)
    return (R + R.transpose(-2, -1)) / 2.0


def make_cm_like_batch(B, n, device, dtype=torch.float32):
    points = torch.randn(B, n, n, device=device, dtype=dtype)
    points = points / (points.norm(dim=-1, keepdim=True) + 1e-8)
    return torch.bmm(points, points.transpose(-2, -1)) * 0.3


def fmt_time(seconds):
    if seconds < 1e-3:
        return f"{seconds*1e6:.1f} us"
    elif seconds < 1.0:
        return f"{seconds*1e3:.2f} ms"
    return f"{seconds:.3f} s"


# ─── Test 0: Newton-Schulz Diagnostic ───

def test_ns_diagnostic(device):
    """Verify Newton-Schulz orthogonalization works on GPU independently."""
    print("\n" + "=" * 70)
    print("  TEST 0: NEWTON-SCHULZ DIAGNOSTIC")
    print("=" * 70)

    for n in [5, 6, 8]:
        B = 1024
        # Create nearly-orthogonal matrix (simulating Jacobi output)
        Q, _ = torch.linalg.qr(torch.randn(B, n, n, device=device))
        # Perturb to ~1e-3 orthogonality error
        noise = torch.randn(B, n, n, device=device) * 1e-3
        V_dirty = Q + noise

        I_n = torch.eye(n, device=device).unsqueeze(0)

        # Before NS
        VtV_before = torch.bmm(V_dirty.transpose(-2, -1), V_dirty)
        orth_before = torch.linalg.norm((VtV_before - I_n).reshape(B, -1), dim=-1).max().item()

        # After NS (2 iterations)
        V_clean = orthogonalize_ns(V_dirty, n_iter=2)
        VtV_after = torch.bmm(V_clean.transpose(-2, -1), V_clean)
        orth_after = torch.linalg.norm((VtV_after - I_n).reshape(B, -1), dim=-1).max().item()

        # After NS (3 iterations for comparison)
        V_clean3 = orthogonalize_ns(V_dirty, n_iter=3)
        VtV_after3 = torch.bmm(V_clean3.transpose(-2, -1), V_clean3)
        orth_after3 = torch.linalg.norm((VtV_after3 - I_n).reshape(B, -1), dim=-1).max().item()

        print(f"  n={n}: before={orth_before:.2e}  "
              f"after(2iter)={orth_after:.2e}  "
              f"after(3iter)={orth_after3:.2e}")

    # Also test with actual Jacobi output
    print(f"\n  --- With actual Jacobi output ---")
    for n in [5, 6]:
        B = 2048
        A = make_symmetric_batch(B, n, device)
        solver = JacobiEigh(max_n=n, max_sweeps=10).to(device)

        # Run Jacobi WITHOUT the NS cleanup
        W = A.clone()
        V = torch.eye(n, device=device).unsqueeze(0).expand(B, -1, -1).clone()
        for _sweep in range(solver.max_sweeps):
            for idx in range(solver._n_pairs):
                p, q = solver._pairs_p[idx], solver._pairs_q[idx]
                app, aqq, apq = W[:, p, p], W[:, q, q], W[:, p, q]
                two_apq = 2.0 * apq
                diff = aqq - app
                abs_2apq = two_apq.abs().clamp(min=1e-30)
                sign_2apq = torch.where(two_apq >= 0,
                    torch.ones_like(two_apq), -torch.ones_like(two_apq))
                tau = diff / (abs_2apq * sign_2apq)
                tau_sign = torch.where(tau >= 0,
                    torch.ones_like(tau), -torch.ones_like(tau))
                t = tau_sign / (tau.abs() + torch.sqrt(1.0 + tau * tau))
                skip = (apq.abs() < 1e-30).float()
                t = t * (1.0 - skip)
                c = 1.0 / torch.sqrt(1.0 + t * t)
                s = t * c
                c_col, s_col = c.unsqueeze(-1), s.unsqueeze(-1)
                Wp = W[:, :, p].clone(); Wq = W[:, :, q].clone()
                W[:, :, p] = c_col * Wp - s_col * Wq
                W[:, :, q] = s_col * Wp + c_col * Wq
                Wp = W[:, p, :].clone(); Wq = W[:, q, :].clone()
                W[:, p, :] = c_col * Wp - s_col * Wq
                W[:, q, :] = s_col * Wp + c_col * Wq
                W[:, p, q] = 0.0; W[:, q, p] = 0.0
                W[:, p, p] = app - t * apq
                W[:, q, q] = aqq + t * apq
                Vp = V[:, :, p].clone(); Vq = V[:, :, q].clone()
                V[:, :, p] = c_col * Vp - s_col * Vq
                V[:, :, q] = s_col * Vp + c_col * Vq

        I_n = torch.eye(n, device=device).unsqueeze(0)
        VtV = torch.bmm(V.transpose(-2, -1), V)
        orth_raw = torch.linalg.norm((VtV - I_n).reshape(B, -1), dim=-1).max().item()

        V_ns = orthogonalize_ns(V, n_iter=2)
        VtV_ns = torch.bmm(V_ns.transpose(-2, -1), V_ns)
        orth_ns = torch.linalg.norm((VtV_ns - I_n).reshape(B, -1), dim=-1).max().item()

        print(f"  Jacobi raw n={n}: orth={orth_raw:.2e}  after NS(2)={orth_ns:.2e}")


# ─── Test 1: Accuracy ───

def test_accuracy(device):
    print("\n" + "=" * 70)
    print("  TEST 1: ACCURACY vs torch.linalg.eigh")
    print("=" * 70)

    validator = EighValidator()
    configs = [
        (3,   4096, "3x3   small"),
        (5,   4096, "5x5   CM matrix size"),
        (6,   4096, "6x6   pentachoron bordered"),
        (8,   2048, "8x8   padded CM"),
        (12,  1024, "12x12 medium"),
        (16,  512,  "16x16 Jacobi boundary"),
    ]

    all_pass = True
    for n, B, label in configs:
        A = make_symmetric_batch(B, n, device)
        ref_vals, ref_vecs = torch.linalg.eigh(A)

        solver = CompiledEigh(max_n=n).to(device)
        our_vals, our_vecs = solver(A)

        val_err = (our_vals - ref_vals).abs().max().item()
        val_mean = (our_vals - ref_vals).abs().mean().item()

        dots = torch.bmm(ref_vecs.transpose(-2, -1), our_vecs)
        alignment = dots.abs().max(dim=-1).values.min().item()

        res_norm, orth_err, max_res = validator(A, our_vals, our_vecs)
        max_orth = orth_err.max().item()

        # Thresholds: eigenval 1e-3, alignment 0.999, orth 1e-4
        ok = val_err < 1e-3 and alignment > 0.999 and max_orth < 1e-4
        if not ok:
            all_pass = False

        print(f"\n  [{'PASS' if ok else 'FAIL'}] {label}  (n={n}, B={B})")
        print(f"    eigenvalue err   max={val_err:.2e}  mean={val_mean:.2e}")
        print(f"    eigvec alignment min={alignment:.8f}")
        print(f"    residual norm    max={max_res.item():.2e}")
        print(f"    orthogonality    max={max_orth:.2e}")

    print(f"\n  --- CM-like spectral distribution ---")
    for n in [5, 6]:
        A = make_cm_like_batch(2048, n, device)
        ref_vals, _ = torch.linalg.eigh(A)
        solver = CompiledEigh(max_n=n).to(device)
        our_vals, our_vecs = solver(A)
        val_err = (our_vals - ref_vals).abs().max().item()
        res_norm, orth_err, max_res = validator(A, our_vals, our_vecs)
        print(f"    CM-like n={n}: val_err={val_err:.2e}  "
              f"res={max_res.item():.2e}  orth={orth_err.max().item():.2e}")

    return all_pass


# ─── Test 2: torch.compile fullgraph ───

def test_compile(device):
    print("\n" + "=" * 70)
    print("  TEST 2: torch.compile(fullgraph=True)")
    print("=" * 70)

    results = {}
    for n, B, label in [(5, 1024, "5x5"), (6, 1024, "6x6"), (8, 512, "8x8")]:
        A = make_symmetric_batch(B, n, device)
        solver = CompiledEigh(max_n=n).to(device)

        try:
            compiled_solver = torch.compile(solver, fullgraph=True)
            vals, vecs = compiled_solver(A)
            sync()
            ref_vals, _ = torch.linalg.eigh(A)
            err = (vals - ref_vals).abs().max().item()
            results[label] = ("PASS", err)
            print(f"  [{label}] fullgraph=True SUCCESS  (val_err={err:.2e})")
        except Exception as e:
            results[label] = ("FAIL", str(e)[:200])
            print(f"  [{label}] COMPILE FAILED: {str(e)[:200]}")

    return all(v[0] == "PASS" for v in results.values())


# ─── Test 3: Throughput ───

def test_benchmark(device):
    print("\n" + "=" * 70)
    print("  TEST 3: GPU THROUGHPUT BENCHMARK")
    print("=" * 70)
    print(f"  Device: {torch.cuda.get_device_name(0)}")
    print(f"  Timing: 10 warmup + 200 repeats\n")

    configs = [
        (5,  1024,  "CM 5x5  B=1024"),
        (5,  4096,  "CM 5x5  B=4096"),
        (5,  8192,  "CM 5x5  B=8192"),
        (6,  1024,  "CM 6x6  B=1024"),
        (6,  4096,  "CM 6x6  B=4096"),
        (6,  8192,  "CM 6x6  B=8192"),
        (8,  2048,  "8x8     B=2048"),
        (16, 1024,  "16x16   B=1024"),
    ]

    print(f"  {'Config':<22} {'eigh ref':>10} {'ours eager':>12} "
          f"{'ours compiled':>14} {'vs ref':>8}")
    print(f"  {'-'*22} {'-'*10} {'-'*12} {'-'*14} {'-'*8}")

    for n, B, label in configs:
        A = make_symmetric_batch(B, n, device)

        ref_time = gpu_timer(lambda: torch.linalg.eigh(A))

        solver = CompiledEigh(max_n=n).to(device)
        eager_time = gpu_timer(lambda: solver(A))

        try:
            compiled_solver = torch.compile(solver, fullgraph=True)
            for _ in range(5):
                compiled_solver(A)
            sync()
            compiled_time = gpu_timer(lambda: compiled_solver(A))
            compiled_str = fmt_time(compiled_time)
            speedup = ref_time / compiled_time
            speedup_str = f"{speedup:.2f}x"
        except Exception:
            compiled_str = "FAIL"
            speedup_str = "N/A"

        print(f"  {label:<22} {fmt_time(ref_time):>10} "
              f"{fmt_time(eager_time):>12} {compiled_str:>14} {speedup_str:>8}")

    print(f"\n  --- High batch stress test ---")
    for n in [5, 6]:
        for B in [16384, 32768]:
            try:
                A = make_symmetric_batch(B, n, device)
                solver = CompiledEigh(max_n=n).to(device)
                compiled_solver = torch.compile(solver, fullgraph=True)
                for _ in range(3):
                    compiled_solver(A)
                sync()
                t = gpu_timer(lambda: compiled_solver(A), warmup=5, repeats=100)
                ref_t = gpu_timer(lambda: torch.linalg.eigh(A), warmup=5, repeats=100)
                print(f"    n={n} B={B}: compiled={fmt_time(t)}  ref={fmt_time(ref_t)}  "
                      f"ratio={ref_t/t:.2f}x  throughput={B/t:.0f}/sec")
            except RuntimeError as e:
                if "out of memory" in str(e).lower():
                    print(f"    n={n} B={B}: OOM")
                    torch.cuda.empty_cache()
                else:
                    raise


# ─── Test 4: Autograd ───

def test_autograd(device):
    print("\n" + "=" * 70)
    print("  TEST 4: AUTOGRAD BACKWARD")
    print("=" * 70)

    for n, B in [(5, 512), (6, 512)]:
        A_ref = make_symmetric_batch(B, n, device).requires_grad_(True)
        vals_ref, vecs_ref = torch.linalg.eigh(A_ref)
        (vals_ref.sum() + (vecs_ref ** 2).sum()).backward()
        grad_ref = A_ref.grad.clone()

        # Eager backward
        A_e = A_ref.detach().clone().requires_grad_(True)
        solver = CompiledEigh(max_n=n).to(device)
        try:
            vals_e, vecs_e = solver(A_e)
            (vals_e.sum() + (vecs_e ** 2).sum()).backward()
            err_e = (A_e.grad - grad_ref).abs().max().item()
            rel_e = err_e / (grad_ref.abs().max().item() + 1e-30)
            print(f"  [{'PASS' if rel_e < 0.1 else 'WARN'}] n={n} eager backward: "
                  f"grad_err={err_e:.2e}  rel={rel_e:.2e}")
        except Exception as e:
            print(f"  [FAIL] n={n} eager backward: {e}")

        # Compiled backward (may break — forward fullgraph is the key win)
        A_c = A_ref.detach().clone().requires_grad_(True)
        try:
            compiled_solver = torch.compile(solver)
            vals_c, vecs_c = compiled_solver(A_c)
            (vals_c.sum() + (vecs_c ** 2).sum()).backward()
            err_c = (A_c.grad - grad_ref).abs().max().item()
            rel_c = err_c / (grad_ref.abs().max().item() + 1e-30)
            print(f"  [{'PASS' if rel_c < 0.1 else 'WARN'}] n={n} compiled backward: "
                  f"grad_err={err_c:.2e}  rel={rel_c:.2e}")
        except Exception as e:
            print(f"  [INFO] n={n} compiled backward: {str(e)[:150]}")
            print(f"         (forward fullgraph is the main win)")


# ─── Test 5: VRAM ───

def test_vram(device):
    print("\n" + "=" * 70)
    print("  TEST 5: VRAM USAGE")
    print("=" * 70)

    for n, B in [(5, 4096), (6, 4096), (6, 8192), (5, 8192)]:
        torch.cuda.empty_cache()
        gc.collect()
        torch.cuda.reset_peak_memory_stats()
        base_mem = torch.cuda.memory_allocated()

        A = make_symmetric_batch(B, n, device)
        solver = CompiledEigh(max_n=n).to(device)
        vals, vecs = solver(A)

        peak_mem = torch.cuda.max_memory_allocated()
        delta_mb = (peak_mem - base_mem) / (1024 ** 2)
        print(f"  n={n} B={B}: peak delta = {delta_mb:.1f} MB")

        del A, solver, vals, vecs
        torch.cuda.empty_cache()
        gc.collect()


# ─── Main ───

def main():
    print("=" * 70)
    print("  CompiledEigh v3 — GPU Benchmark Suite")
    print("=" * 70)

    if not torch.cuda.is_available():
        print("\n  No CUDA. Run on Colab with A100/H100.")
        sys.exit(1)

    device = torch.device('cuda')
    print(f"\n  GPU:     {torch.cuda.get_device_name(0)}")
    print(f"  CUDA:    {torch.version.cuda}")
    print(f"  PyTorch: {torch.__version__}")
    mem_gb = torch.cuda.get_device_properties(0).total_memory / (1024**3)
    print(f"  VRAM:    {mem_gb:.1f} GB")
    print(f"  TF32 matmul: {torch.backends.cuda.matmul.allow_tf32}")
    print(f"  float32 precision: {torch.get_float32_matmul_precision()}")

    test_ns_diagnostic(device)
    acc_ok = test_accuracy(device)
    compile_ok = test_compile(device)
    test_benchmark(device)
    test_autograd(device)
    test_vram(device)

    print("\n" + "=" * 70)
    print("  SUMMARY")
    print("=" * 70)
    print(f"  Accuracy:  {'PASS' if acc_ok else 'FAIL'}")
    print(f"  Compile:   {'PASS' if compile_ok else 'FAIL'}")
    print("=" * 70)


if __name__ == '__main__':
    main()