perf(solvers): route eager Davidson RR GEMMs through OpenBLAS zgemm (measured 1.05x serial / 1.02x k-parallel)

[wladerer]

What

Routes the eager davidson_batched Rayleigh-Ritz GEMMs — the subspace build S = V·HV^H, the Ritz combines X = U^T·V / HX = U^T·HV, and the two-pass projection in _orthonormalize_b — through OpenBLAS cblas_zgemm via a tiny ctypes-loaded native symbol, because torch's CPU complex128 matmul is measured 2.6-5x slower per-op than OpenBLAS zgemm on exactly these batched shapes (asus: 210 vs 718 µs @ batch 8 × 754×754; 1237 vs 6354 µs @ 27×1260×1260). FFTs are already at torch/FFTW parity and are untouched.

• src/gradwave/solvers/native/blas_gemm.c — one blas_zbmm symbol (strided batch loop over cblas_zgemm), compiled into the same libdavnative.so as the native Davidson solver (one .so, two symbols; scripts/build_native_solver.sh extended, CI builds it via the same script).
• src/gradwave/solvers/blas_routing.py — zbmm(a, b, *, conj_b_t, t_a): native when the .so loads and operands are complex128 CPU batched tensors in a GEMM-expressible layout; silent torch fallback otherwise (identical-math micro-routing, unlike davidson-native's hard error), with native/fallback dispatch counters so a bench can prove engagement.
• Env gate GRADWAVE_BLAS_GEMM ∈ {auto (default), on, off} — read per call (in-process A/B safe, no import-frozen env); on errors if the .so is missing, off is pure torch.
• Layout lesson the first A/B run taught (caught by the engagement counter, not the wall clock): the real Davidson operands are a transpose view (_orthonormalize_b returns q.transpose(-1,-2)) and a u[:, :, :nb] slice — a strict is_contiguous() gate silently rejects them all. GEMM's lda/stride arguments express both (slices via lda = stride(-2), transpose views via a trans-flag flip); only ConjTrans-of-transposed-storage still falls back (one small GEMM per solve).
• OpenBLAS threading follows torch.get_num_threads(), so the k-parallel thread pool (torch pinned to 1 inside tasks) can't oversubscribe workers × cores.

Numerically equivalent by construction: same contraction, different kernel (last-bit BLAS rounding only). The eigensolve runs under no_grad; zbmm additionally refuses grad-tracked operands so autograd can never be silently detached.

Measured (asus, 22-core, Al-4 conv cell 4³ fermi-dirac, ecut=30 Ry, PBE, ONCV — the campaign standard; best of 2 per arm, idle box under the bench-lock protocol, engagement proven by the dispatch counters printed next to each wall time)

arm	off	on	speedup
eager serial	36.03 s	34.40 s	1.047×
k_parallel=8	11.68 s	11.44 s	1.021×

Identical free energy to every printed digit (F = −7506.90978722 eV in all 8 runs) and identical SCF iteration counts (10) in every arm. n_fallback = 0 in both on arms.

Honest read: the 2.6-5× per-op deficit shrinks to ~4.7% end-to-end (serial) because the RR GEMMs are a modest fraction of eager wall (Amdahl), and to ~2% under k_parallel=8 where the per-task GEMMs are smaller and BLAS runs serial inside the pool. A real but small win that costs nothing when the .so is absent, and composes with the k-parallel path.

Validation note: VASP is not runnable in this environment; QE is the external correctness proxy for the SCF this rides on (unchanged here — same eigenpairs, same iterations, kernel-level substitution only).

Tests

tests/unit/test_blas_routing.py — zbmm vs torch.matmul on random complex128 batches for all routed trans/conj combos (rtol 1e-12), the view layouts davidson actually produces (transpose view, last-dim slice) asserted to route natively, out-of-scope layouts asserted to fall back correctly, env gating (off bit-identical torch, on without .so raises, auto silent fallback), and a davidson_batched A/B on a random Hermitian operator: same eigenvalues to 1e-10, same n_iter, routing on vs off. Skips cleanly when the .so is absent.

Fast gate green on asus; ruff / ty / lint-imports / doc-refs clean; branch on current origin/main.

🤖 Generated with Claude Code

https://claude.ai/code/session_018j6BBnPtBSjkrnASSc52mT

[wladerer]

Parked by measured verdict: exact and well-tested, but 1.047× serial / 1.021× k-parallel e2e (Al-4, asus, engagement-verified) doesn't justify a new C surface + routing layer — the per-op 2.6-5× OpenBLAS zgemm advantage Amdahl-shrinks to ~5% because the RR GEMMs are a thin slice of eager wall, and the main path now runs davidson-native (7.5s Al-4). Branch kept for the record; reopen only if the eager fallback scopes (USPP/mGGA/Fock) become a measured bottleneck at larger nb.