Create eigh_cuda_kernel.py

eigh_cuda_kernel.py

@@ -0,0 +1,397 @@
+"""
+fl_eigh_cuda.py — CUDA FL Hybrid Eigh via CuPy RawKernel.
+
+Compiles at runtime using NVRTC (part of CUDA toolkit, already installed).
+No ninja, no C++ compiler, no build system. Just pip install cupy-cuda12x.
+
+PyTorch <-> CuPy via DLPack (zero-copy).
+
+Usage:
+    from fl_eigh_cuda import fl_eigh_cuda
+    evals, evecs = fl_eigh_cuda(A)  # A is [B, 6, 6] PyTorch CUDA tensor
+"""
+
+import math, time, gc, sys
+import torch
+from torch import Tensor
+from typing import Tuple
+
+torch.backends.cuda.matmul.allow_tf32 = False
+torch.backends.cudnn.allow_tf32 = False
+torch.set_float32_matmul_precision('highest')
+
+
+_KERNEL_SRC = r"""
+extern "C" __global__
+void fl_eigh_kernel(
+    const float* __restrict__ A_in,
+    float* __restrict__ evals_out,
+    float* __restrict__ evecs_out,
+    int B
+) {
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+    if (tid >= B) return;
+
+    const int NN = 6;
+    const int N2 = 36;
+
+    // Load A and pre-scale
+    double a[36];
+    double frob_sq = 0.0;
+    for (int i = 0; i < N2; i++) {
+        a[i] = (double)A_in[tid * N2 + i];
+        frob_sq += a[i] * a[i];
+    }
+    double scale = sqrt(frob_sq / 6.0);
+    if (scale < 1e-12) scale = 1e-12;
+    double inv_s = 1.0 / scale;
+    for (int i = 0; i < N2; i++) a[i] *= inv_s;
+
+    // Phase 1: FL coefficients (fp64)
+    double c[7];
+    for (int i = 0; i < 7; i++) c[i] = 0.0;
+    c[6] = 1.0;
+
+    double m[36];
+    for (int i = 0; i < N2; i++) m[i] = 0.0;
+
+    for (int k = 1; k <= NN; k++) {
+        double mn[36];
+        for (int i = 0; i < NN; i++) {
+            for (int j = 0; j < NN; j++) {
+                double acc = 0.0;
+                for (int l = 0; l < NN; l++)
+                    acc += a[i*NN+l] * m[l*NN+j];
+                if (i == j) acc += c[NN-k+1];
+                mn[i*NN+j] = acc;
+            }
+        }
+        double tr = 0.0;
+        for (int i = 0; i < NN; i++)
+            for (int l = 0; l < NN; l++)
+                tr += a[i*NN+l] * mn[l*NN+i];
+        c[NN-k] = -tr / (double)k;
+        for (int i = 0; i < N2; i++) m[i] = mn[i];
+    }
+
+    // Phase 2: Laguerre + deflation + polish (fp64)
+    double diag[6];
+    for (int i = 0; i < NN; i++) diag[i] = a[i*NN+i];
+    for (int pass = 0; pass < NN-1; pass++)
+        for (int j = 0; j < NN-1; j++)
+            if (diag[j] > diag[j+1]) {
+                double tmp = diag[j]; diag[j] = diag[j+1]; diag[j+1] = tmp;
+            }
+    for (int i = 0; i < NN; i++)
+        diag[i] += -1e-4 + 2e-4 * (double)i / 5.0;
+
+    double cl[7];
+    for (int i = 0; i < 7; i++) cl[i] = c[i];
+
+    double roots[6];
+
+    for (int ri = 0; ri < NN; ri++) {
+        int deg = NN - ri;
+        double z = diag[ri];
+        for (int lag = 0; lag < 5; lag++) {
+            double pv = cl[deg], dp = 0.0, d2 = 0.0;
+            for (int j = deg - 1; j >= 0; j--) {
+                d2 = d2 * z + dp;
+                dp = dp * z + pv;
+                pv = pv * z + cl[j];
+            }
+            if (fabs(pv) > 1e-30) {
+                double G = dp / pv;
+                double H = G * G - 2.0 * d2 / pv;
+                double disc = ((double)(deg-1)) * ((double)deg * H - G * G);
+                if (disc < 0.0) disc = 0.0;
+                double sq = sqrt(disc);
+                double gp = G + sq, gm = G - sq;
+                double den = (fabs(gp) >= fabs(gm)) ? gp : gm;
+                if (fabs(den) > 1e-20)
+                    z -= (double)deg / den;
+            }
+        }
+        roots[ri] = z;
+        if (deg > 1) {
+            double b = cl[deg];
+            for (int j = deg - 1; j > 0; j--) {
+                double bn = cl[j] + z * b;
+                cl[j] = b;
+                b = bn;
+            }
+            cl[0] = b;
+        }
+    }
+
+    // Newton polish
+    for (int pol = 0; pol < 3; pol++)
+        for (int ri = 0; ri < NN; ri++) {
+            double pv = c[NN], dp = 0.0;
+            for (int j = NN - 1; j >= 0; j--) {
+                dp = dp * roots[ri] + pv;
+                pv = pv * roots[ri] + c[j];
+            }
+            if (fabs(dp) > 1e-30)
+                roots[ri] -= pv / dp;
+        }
+
+    // Phase 3: Eigenvectors via interleaved FL+Horner (fp64)
+    float evecs[36];
+
+    for (int ei = 0; ei < NN; ei++) {
+        double lam = roots[ei];
+        double m_loc[36], r_loc[36];
+        for (int i = 0; i < N2; i++) m_loc[i] = 0.0;
+
+        for (int k = 1; k <= NN; k++) {
+            double mn_loc[36];
+            for (int i = 0; i < NN; i++)
+                for (int j = 0; j < NN; j++) {
+                    double acc = 0.0;
+                    for (int l = 0; l < NN; l++)
+                        acc += a[i*NN+l] * m_loc[l*NN+j];
+                    if (i == j) acc += c[NN-k+1];
+                    mn_loc[i*NN+j] = acc;
+                }
+            if (k == 1)
+                for (int i = 0; i < N2; i++) r_loc[i] = mn_loc[i];
+            else
+                for (int i = 0; i < N2; i++) r_loc[i] = r_loc[i] * lam + mn_loc[i];
+            for (int i = 0; i < N2; i++) m_loc[i] = mn_loc[i];
+        }
+
+        int best_j = 0;
+        double best_norm = -1.0;
+        for (int j = 0; j < NN; j++) {
+            double col_sq = 0.0;
+            for (int i = 0; i < NN; i++)
+                col_sq += r_loc[i*NN+j] * r_loc[i*NN+j];
+            if (col_sq > best_norm) { best_norm = col_sq; best_j = j; }
+        }
+        double vnorm = 0.0;
+        double vec[6];
+        for (int i = 0; i < NN; i++) {
+            vec[i] = r_loc[i*NN + best_j];
+            vnorm += vec[i] * vec[i];
+        }
+        vnorm = sqrt(vnorm) + 1e-30;
+        for (int i = 0; i < NN; i++)
+            evecs[i*NN + ei] = (float)(vec[i] / vnorm);
+    }
+
+    // Phase 4: Newton-Schulz (fp32, 2 iters)
+    for (int ns = 0; ns < 2; ns++) {
+        float y[36], t_m[36], vn[36];
+        for (int i = 0; i < NN; i++)
+            for (int j = 0; j < NN; j++) {
+                float acc = 0.0f;
+                for (int l = 0; l < NN; l++)
+                    acc += evecs[l*NN+i] * evecs[l*NN+j];
+                y[i*NN+j] = acc;
+            }
+        for (int i = 0; i < NN; i++)
+            for (int j = 0; j < NN; j++)
+                t_m[i*NN+j] = ((i==j) ? 3.0f : 0.0f) - y[i*NN+j];
+        for (int i = 0; i < NN; i++)
+            for (int j = 0; j < NN; j++) {
+                float acc = 0.0f;
+                for (int l = 0; l < NN; l++)
+                    acc += evecs[i*NN+l] * t_m[l*NN+j];
+                vn[i*NN+j] = 0.5f * acc;
+            }
+        for (int i = 0; i < N2; i++) evecs[i] = vn[i];
+    }
+
+    // Phase 5: Rayleigh quotient (fp32)
+    float af[36];
+    for (int i = 0; i < N2; i++) af[i] = (float)a[i];
+
+    float evals_local[6];
+    for (int ei = 0; ei < NN; ei++) {
+        float lam_f = 0.0f;
+        for (int l = 0; l < NN; l++) {
+            float av = 0.0f;
+            for (int mm = 0; mm < NN; mm++)
+                av += af[l*NN+mm] * evecs[mm*NN+ei];
+            lam_f += evecs[l*NN+ei] * av;
+        }
+        evals_local[ei] = lam_f * (float)scale;
+    }
+
+    // Sort ascending + permute
+    int perm[6];
+    for (int i = 0; i < NN; i++) perm[i] = i;
+    for (int pass = 0; pass < NN-1; pass++)
+        for (int j = 0; j < NN-1; j++)
+            if (evals_local[j] > evals_local[j+1]) {
+                float tmp = evals_local[j]; evals_local[j] = evals_local[j+1]; evals_local[j+1] = tmp;
+                int ptmp = perm[j]; perm[j] = perm[j+1]; perm[j+1] = ptmp;
+            }
+
+    for (int i = 0; i < NN; i++)
+        evals_out[tid * NN + i] = evals_local[i];
+    for (int j_out = 0; j_out < NN; j_out++) {
+        int j_src = perm[j_out];
+        for (int i = 0; i < NN; i++)
+            evecs_out[tid * N2 + i*NN + j_out] = evecs[i*NN + j_src];
+    }
+}
+"""
+
+# ═══════════════════════════════════════════════════════════════════════
+# CuPy compilation + PyTorch wrapper
+# ═══════════════════════════════════════════════════════════════════════
+
+_kernel = None
+
+def _get_kernel():
+    global _kernel
+    if _kernel is not None:
+        return _kernel
+    import cupy
+    print("  Compiling via NVRTC...", end=" ", flush=True)
+    _kernel = cupy.RawKernel(_KERNEL_SRC, 'fl_eigh_kernel')
+    # Force compilation now (not on first launch)
+    _kernel.compile()
+    print("done.")
+    return _kernel
+
+
+def fl_eigh_cuda(A: Tensor) -> Tuple[Tensor, Tensor]:
+    """CUDA FL Hybrid Eigendecomposition for [B, 6, 6] symmetric matrices.
+
+    Uses CuPy RawKernel (NVRTC). Zero-copy PyTorch interop via data_ptr.
+    """
+    assert A.is_cuda and A.shape[-2:] == (6, 6), f"Need CUDA [B,6,6], got {A.shape}"
+    B = A.shape[0]
+    kernel = _get_kernel()
+
+    A_contig = A.contiguous().float()
+    evals = torch.empty(B, 6, device=A.device, dtype=torch.float32)
+    evecs = torch.empty(B, 6, 6, device=A.device, dtype=torch.float32)
+
+    import cupy
+    # Raw pointers — zero copy, no DLPack needed
+    a_ptr = cupy.cuda.MemoryPointer(
+        cupy.cuda.UnownedMemory(A_contig.data_ptr(), A_contig.nelement() * 4, None), 0)
+    ev_ptr = cupy.cuda.MemoryPointer(
+        cupy.cuda.UnownedMemory(evals.data_ptr(), evals.nelement() * 4, None), 0)
+    vc_ptr = cupy.cuda.MemoryPointer(
+        cupy.cuda.UnownedMemory(evecs.data_ptr(), evecs.nelement() * 4, None), 0)
+
+    threads = 128
+    blocks = (B + threads - 1) // threads
+
+    # Launch on PyTorch's current CUDA stream
+    stream = cupy.cuda.ExternalStream(torch.cuda.current_stream().cuda_stream)
+    with stream:
+        kernel((blocks,), (threads,),
+               (a_ptr, ev_ptr, vc_ptr, B))
+
+    return evals, evecs
+
+
+# ═══════════════════════════════════════════════════════════════════════
+# Math purity test
+# ═══════════════════════════════════════════════════════════════════════
+
+def math_test(A, vals, vecs):
+    B,n,_=A.shape; dev=A.device
+    Ad=A.double(); vd=vals.double(); Vd=vecs.double()
+    AV=torch.bmm(Ad,Vd); VL=Vd*vd.unsqueeze(-2)
+    An=Ad.reshape(B,-1).norm(dim=-1,keepdim=True).clamp(min=1e-30)
+    res=(AV-VL).norm(dim=-2)/An
+    VtV=torch.bmm(Vd.mT,Vd); I=torch.eye(n,device=dev,dtype=torch.float64).unsqueeze(0)
+    orth=(VtV-I).reshape(B,-1).norm(dim=-1)
+    recon=torch.bmm(Vd*vd.unsqueeze(-2),Vd.mT)
+    recon_err=(Ad-recon).reshape(B,-1).norm(dim=-1)/An.squeeze(-1)
+    tr_err=(Ad.diagonal(dim1=-2,dim2=-1).sum(-1)-vd.sum(-1)).abs()
+    det_A=torch.linalg.det(Ad); det_err=(det_A-vd.prod(-1)).abs()/det_A.abs().clamp(min=1e-30)
+    return dict(res_max=res.max().item(), res_mean=res.mean().item(),
+                orth_max=orth.max().item(), orth_mean=orth.mean().item(),
+                recon_max=recon_err.max().item(), recon_mean=recon_err.mean().item(),
+                tr_max=tr_err.max().item(), det_max=det_err.max().item())
+
+
+# ═══════════════════════════════════════════════════════════════════════
+# Benchmark
+# ═══════════════════════════════════════════════════════════════════════
+
+def sync(): torch.cuda.synchronize()
+def gt(fn,w=20,r=200):
+    for _ in range(w): fn()
+    sync(); t=time.perf_counter()
+    for _ in range(r): fn()
+    sync(); return (time.perf_counter()-t)/r
+def fmt(s):
+    if s<1e-3: return f"{s*1e6:.1f}us"
+    if s<1: return f"{s*1e3:.2f}ms"
+    return f"{s:.3f}s"
+
+
+def main():
+    if not torch.cuda.is_available(): sys.exit(1)
+    dev=torch.device('cuda')
+    p=torch.cuda.get_device_properties(0)
+    print("="*72)
+    print("  FL Eigh CUDA Kernel (CuPy/NVRTC)")
+    print("="*72)
+    print(f"  {p.name}")
+    print(f"  PyTorch {torch.__version__}")
+
+    N=6; B=4096
+    A=(lambda R:(R+R.mT)/2)(torch.randn(B,N,N,device=dev))
+    rv,rV=torch.linalg.eigh(A)
+
+    _get_kernel()
+
+    # Accuracy
+    print(f"\n  ACCURACY (n={N} B={B})")
+    cv,cV=fl_eigh_cuda(A)
+    ve=(cv-rv).abs().max().item()
+    dots=torch.bmm(rV.double().mT,cV.double()).abs().max(dim=-1).values.min().item()
+    print(f"  CUDA FL: val={ve:.1e} align={dots:.6f}")
+
+    # Math purity
+    mc=math_test(A,rv,rV); mf=math_test(A,cv,cV)
+    wins=0
+    print(f"\n  MATH PURITY: CUDA FL vs cuSOLVER")
+    print(f"  {'Property':<28} {'cuSOLVER':>10} {'CUDA FL':>10} {'Win':>6}")
+    for key in ['res_max','res_mean','orth_max','orth_mean','recon_max','recon_mean','tr_max','det_max']:
+        vc=mc[key]; vf=mf[key]; w='FL' if vf<vc else 'cuS'
+        if vf<vc: wins+=1
+        print(f"  {key:<28} {vc:>10.1e} {vf:>10.1e} {w:>6}")
+    print(f"\n  CUDA FL wins {wins}/8")
+
+    # Throughput
+    print(f"\n  THROUGHPUT (n={N} B={B})")
+    tr=gt(lambda:torch.linalg.eigh(A))
+    tc=gt(lambda:fl_eigh_cuda(A))
+    print(f"  cuSOLVER:  {fmt(tr)}")
+    print(f"  CUDA FL:   {fmt(tc)} ({tr/tc:.2f}x)")
+
+    # Batch scaling
+    print(f"\n  BATCH SCALING (n={N})")
+    print(f"  {'B':>6}  {'cuSOLVER':>10}  {'CUDA FL':>10}  {'ratio':>7}")
+    for Bx in [256,512,1024,2048,4096,8192,16384,32768]:
+        try:
+            Ax=(lambda R:(R+R.mT)/2)(torch.randn(Bx,N,N,device=dev))
+            t1=gt(lambda:torch.linalg.eigh(Ax),10,100)
+            t2=gt(lambda:fl_eigh_cuda(Ax),10,100)
+            print(f"  {Bx:>6}  {fmt(t1):>10}  {fmt(t2):>10}  {t1/t2:>6.2f}x")
+            del Ax
+        except RuntimeError:
+            print(f"  {Bx:>6}  OOM"); torch.cuda.empty_cache()
+
+    # Memory
+    print(f"\n  MEMORY (n={N} B={B})")
+    for lbl,fn in [("cuSOLVER",lambda:torch.linalg.eigh(A)),("CUDA FL",lambda:fl_eigh_cuda(A))]:
+        torch.cuda.empty_cache(); gc.collect(); torch.cuda.reset_peak_memory_stats()
+        base=torch.cuda.memory_allocated(); fn(); sync()
+        print(f"  {lbl:<12} {(torch.cuda.max_memory_allocated()-base)/1024**2:.1f}MB")
+
+    print("="*72)
+
+
+if __name__=='__main__': main()