""" 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()