v0.8.0: SSOR preconditioner for SPD systems (#28)

Add ssor_preconditioner(A, omega=1.0) to trnsolver.iterative. Applies
M^{-1} r via forward triangular solve (D + ωL) t = r, diagonal scaling
v = ω(2-ω) diag(A) ⊙ t, and backward solve (D + ωL^T) z = v. ω=1 is
symmetric Gauss-Seidel; converges faster than Jacobi on coupled matrices
(1D Laplacian, FEM stiffness). BF16/FP16 promoted to FP32 in factory and
closure. SSOR benchmark added to bench_solver.py. Closes SSOR item in #28.

Author: scttfrdmn

CHANGELOG.md

@@ -7,6 +7,23 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+## [0.8.0] — 2026-04-17
+
+### Added
+
+- **`ssor_preconditioner(A, omega=1.0)`** in `trnsolver.iterative`. SSOR
+  (Symmetric Successive Over-Relaxation) preconditioner for symmetric positive
+  definite A. Applies M^{-1} r via three steps: (1) forward triangular solve
+  (D + ω L) t = r, (2) diagonal scaling v = ω(2-ω) diag(A) ⊙ t, (3) backward
+  triangular solve (D + ω L^T) z = v. At ω = 1 this reduces to the symmetric
+  Gauss-Seidel step. M_SSOR is SPD when A is SPD and ω ∈ (0, 2), so it is safe
+  as a CG preconditioner. SSOR outperforms scalar Jacobi on matrices with
+  off-diagonal coupling (1D Laplacian, FEM stiffness, density-fitting metrics).
+  BF16/FP16 inputs are promoted to FP32 in both the factory and the returned
+  closure (consistent with the dtype-promotion pattern in #19). Adds the final
+  preconditioner item from the Phase 3 tracker (#28); SSOR benchmark added to
+  `benchmarks/bench_solver.py`.
+
 ## [0.7.0] — 2026-04-17
 
 ### Added

README.md

@@ -57,7 +57,8 @@ Ap = trnsolver.pinv(A)
 # Iterative solvers with preconditioners
 precond = trnsolver.jacobi_preconditioner(A)
 blk_precond = trnsolver.block_jacobi_preconditioner(A, block_size=16)
-x, iters, res = trnsolver.cg(A, b, M=blk_precond, tol=1e-8)
+ssor_precond = trnsolver.ssor_preconditioner(A, omega=1.0)  # symmetric Gauss-Seidel
+x, iters, res = trnsolver.cg(A, b, M=ssor_precond, tol=1e-8)
 x, iters, res = trnsolver.gmres(A, b, tol=1e-6)
 ```
 
@@ -76,6 +77,8 @@ Demonstrates the self-consistent-field iteration: build Fock matrix → solve ge
 
 ## Status
 
+**v0.8.0** — `ssor_preconditioner(A, omega=1.0)` (#28). SSOR preconditioner for SPD systems: two triangular solves + diagonal scaling per application. Outperforms scalar Jacobi on coupled matrices (1D Laplacian, FEM stiffness). ω ∈ (0, 2); ω=1 is symmetric Gauss-Seidel.
+
 **v0.7.0** — BF16/FP16 dtype support across the full public API (#19). All entry points (`cholesky`, `lu`, `qr`, `solve`, `solve_spd`, `inv_spd`, `pinv`, `inv_sqrt_spd`, `inv_sqrt_spd_ns`, `eigh`, `eigh_generalized`, `cg`, `gmres`, `block_jacobi_preconditioner`) accept BF16/FP16 inputs, upcast to FP32 internally, and restore the original dtype on output.
 
 **v0.6.0** — `eigh` subspace rotation refinement: one Rayleigh-Ritz step (V^T A V re-diagonalization) after Householder-QR reduces eigenvector residuals by 1–2 orders of magnitude for n ≥ 64 (#31). `solve_spd` gains `iterative_refinement=True`: mixed-precision FP64 residual + second Cholesky solve for SPD systems with cond up to ~1e7 (#32).
@@ -87,16 +90,17 @@ Demonstrates the self-consistent-field iteration: build Fock matrix → solve ge
 | Eigensolvers | `eigh`, `eigh_generalized` | `svd` (Jacobi-SVD, Phase 3) |
 | Factorizations | `cholesky`, `lu`, `qr`, `pinv` | `schur` (implicit-shift QR, Phase 3) |
 | Direct solvers | `solve`, `solve_spd`, `inv_spd`, `inv_sqrt_spd`, `inv_sqrt_spd_ns` | — |
-| Iterative | `cg` (w/ preconditioner), `gmres` | SSOR (#16, v0.6.0) |
-| Preconditioners | `jacobi_preconditioner`, `block_jacobi_preconditioner` | SSOR (v0.6.0) |
+| Iterative | `cg` (w/ preconditioner), `gmres` | — |
+| Preconditioners | `jacobi_preconditioner`, `block_jacobi_preconditioner`, `ssor_preconditioner` | — |
 
 **Roadmap:**
 - **v0.4.0** — NKI Householder-QR `eigh` validated on trn1.2xlarge (#9, #12, #38)
 - **v0.4.1** — `eigh_generalized` NKI triangular-solve path via `trnblas.trsm` (#11)
 - **v0.5.0** — Newton-Schulz trnblas.gemm (#14), FP64 CG/GMRES dots + Rayleigh refinement (#27), block-Jacobi (#16), `pinv` (#22)
 - **v0.6.0** — eigh subspace rotation refinement (#31), solve_spd iterative refinement (#32)
 - **v0.7.0** — BF16/FP16 across the full API (#19) ✓
-- **v0.8.0+** — SSOR preconditioner, multi-NeuronCore parallel Jacobi (#20)
+- **v0.8.0** — SSOR preconditioner (#28) ✓
+- **v0.9.0+** — multi-NeuronCore parallel Jacobi (#20)
 
 ## Operations
 
@@ -117,6 +121,7 @@ Demonstrates the self-consistent-field iteration: build Fock matrix → solve ge
 | Iterative | `gmres` | GMRES (general systems) |
 | Iterative | `jacobi_preconditioner` | Diagonal preconditioner for CG |
 | Iterative | `block_jacobi_preconditioner` | Block-diagonal Cholesky preconditioner for CG |
+| Iterative | `ssor_preconditioner` | SSOR / symmetric Gauss-Seidel preconditioner for CG |
 
 ## Benchmarks
 
@@ -132,7 +137,7 @@ All six siblings are on PyPI, along with the umbrella meta-package:
 | [trnfft](https://github.com/trnsci/trnfft) | FFT and complex-valued tensors | v0.8.0 |
 | [trnblas](https://github.com/trnsci/trnblas) | BLAS Level 1–3 | v0.4.0 |
 | [trnrand](https://github.com/trnsci/trnrand) | Philox / Sobol / Halton RNG | v0.1.0 |
-| trnsolver | Linear solvers and eigendecomposition | **v0.7.0** |
+| trnsolver | Linear solvers and eigendecomposition | **v0.8.0** |
 | [trnsparse](https://github.com/trnsci/trnsparse) | Sparse matrix operations | v0.1.1 |
 | [trntensor](https://github.com/trnsci/trntensor) | Tensor contractions (einsum, TT/Tucker) | v0.1.1 |
 

benchmarks/bench_solver.py

@@ -135,6 +135,10 @@ def test_cg_jacobi_precond(self, benchmark, spd_matrix, random_vector):
         M = trnsolver.jacobi_preconditioner(spd_matrix)
         benchmark(trnsolver.cg, spd_matrix, random_vector, M=M)
 
+    def test_cg_ssor_precond(self, benchmark, spd_matrix, random_vector):
+        M = trnsolver.ssor_preconditioner(spd_matrix)
+        benchmark(trnsolver.cg, spd_matrix, random_vector, M=M)
+
 
 class TestGMRES:
     def test_gmres(self, benchmark, random_matrix, random_vector):

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "trnsolver"
-version = "0.7.0"
+version = "0.8.0"
 description = "Linear solvers and eigendecomposition for AWS Trainium via NKI"
 readme = "README.md"
 license = "Apache-2.0"

tests/test_dtype.py

@@ -162,3 +162,13 @@ def test_gmres(self, dtype):
         assert x.dtype == dtype
         x_ref, _, _ = trnsolver.gmres(A_fp32, b_fp32, tol=1e-4)
         assert torch.allclose(x.float(), x_ref, atol=2e-1, rtol=1e-1)
+
+    def test_ssor(self, dtype):
+        A_fp32 = _make_spd(16)
+        b_fp32 = torch.randn(16)
+        A = A_fp32.to(dtype)
+        b = b_fp32.to(dtype)
+        M = trnsolver.ssor_preconditioner(A)
+        x, iters, res = trnsolver.cg(A, b, M=M, tol=1e-4)
+        assert x.dtype == dtype
+        assert iters > 0

tests/test_iterative.py

@@ -115,6 +115,72 @@ def test_block_size_1_matches_jacobi(self, spd_matrix):
         np.testing.assert_allclose(M_block1(r).numpy(), M_scalar(r).numpy(), atol=1e-6)
 
 
+class TestSSORPreconditioner:
+    def test_reduces_iterations_vs_jacobi(self):
+        # 1D Laplacian — classical test case where symmetric Gauss-Seidel (omega=1)
+        # beats scalar Jacobi in CG iteration count.
+        n = 64
+        ones = torch.ones(n - 1)
+        A = 2 * torch.eye(n) - torch.diag(ones, 1) - torch.diag(ones, -1)
+        b = torch.randn(n)
+
+        M_jac = trnsolver.jacobi_preconditioner(A)
+        M_ssor = trnsolver.ssor_preconditioner(A, omega=1.0)
+
+        _, iters_jac, _ = trnsolver.cg(A, b, tol=1e-8, maxiter=500, M=M_jac)
+        _, iters_ssor, _ = trnsolver.cg(A, b, tol=1e-8, maxiter=500, M=M_ssor)
+
+        assert iters_ssor < iters_jac
+
+    def test_omega_1_converges(self, spd_matrix):
+        n = 16
+        A = spd_matrix(n)
+        b = torch.randn(n)
+        M = trnsolver.ssor_preconditioner(A, omega=1.0)
+        x, iters, res = trnsolver.cg(A, b, tol=1e-8, maxiter=200, M=M)
+        assert res < 1e-7
+        assert iters > 0
+
+    def test_omega_range(self, spd_matrix):
+        n = 16
+        A = spd_matrix(n)
+        b = torch.randn(n)
+        for omega in (0.5, 1.5):
+            M = trnsolver.ssor_preconditioner(A, omega=omega)
+            _, iters, res = trnsolver.cg(A, b, tol=1e-6, maxiter=200, M=M)
+            assert res < 1e-5, f"omega={omega} did not converge: res={res}"
+
+    def test_invalid_omega(self, spd_matrix):
+        A = spd_matrix(8)
+        with pytest.raises(ValueError, match="omega"):
+            trnsolver.ssor_preconditioner(A, omega=0.0)
+        with pytest.raises(ValueError, match="omega"):
+            trnsolver.ssor_preconditioner(A, omega=2.0)
+
+    def test_apply_matches_manual(self):
+        # Small n=4 case: manually compute forward/backward solve and compare.
+        torch.manual_seed(5)
+        n = 4
+        Q, _ = torch.linalg.qr(torch.randn(n, n))
+        A = Q @ torch.diag(torch.tensor([1.0, 2.0, 3.0, 4.0])) @ Q.T
+        r = torch.randn(n)
+        omega = 1.0
+
+        d = torch.diagonal(A)
+        L_factor = omega * torch.tril(A, diagonal=-1) + torch.diag(d)
+        U_factor = L_factor.T
+        dscale = omega * (2.0 - omega) * d
+
+        t = torch.linalg.solve_triangular(L_factor, r.unsqueeze(-1), upper=False).squeeze(-1)
+        v = dscale * t
+        z_ref = torch.linalg.solve_triangular(U_factor, v.unsqueeze(-1), upper=True).squeeze(-1)
+
+        M = trnsolver.ssor_preconditioner(A, omega=omega)
+        z = M(r)
+
+        np.testing.assert_allclose(z.numpy(), z_ref.numpy(), atol=1e-6)
+
+
 class TestGMRES:
     def test_identity(self):
         A = torch.eye(4)

trnsolver/__init__.py

@@ -6,7 +6,7 @@
 Part of the trnsci scientific computing suite.
 """
 
-__version__ = "0.7.0"
+__version__ = "0.8.0"
 
 # Eigenvalue decomposition
 from .eigen import eigh, eigh_generalized
@@ -15,7 +15,13 @@
 from .factor import cholesky, inv_spd, inv_sqrt_spd, inv_sqrt_spd_ns, lu, pinv, qr, solve, solve_spd
 
 # Iterative solvers
-from .iterative import block_jacobi_preconditioner, cg, gmres, jacobi_preconditioner
+from .iterative import (
+    block_jacobi_preconditioner,
+    cg,
+    gmres,
+    jacobi_preconditioner,
+    ssor_preconditioner,
+)
 
 # Backend control
 from .nki import HAS_NKI, get_backend, set_backend
@@ -39,6 +45,7 @@
     "gmres",
     "jacobi_preconditioner",
     "block_jacobi_preconditioner",
+    "ssor_preconditioner",
     # Backend
     "HAS_NKI",
     "set_backend",

trnsolver/iterative.py

@@ -82,6 +82,56 @@ def apply(r: torch.Tensor) -> torch.Tensor:
     return apply
 
 
+def ssor_preconditioner(
+    A: torch.Tensor,
+    omega: float = 1.0,
+) -> Callable[[torch.Tensor], torch.Tensor]:
+    """Build an SSOR preconditioner for symmetric positive definite A.
+
+    For symmetric A = D + L + L^T (D diagonal, L strictly lower triangular),
+    the SSOR preconditioner with relaxation ω ∈ (0, 2) applies M^{-1} r via:
+      1. Forward solve:  (D + ω L) t = r
+      2. Scale:          v = ω(2-ω) * diag(A) ⊙ t
+      3. Backward solve: (D + ω L^T) z = v
+
+    At ω = 1 this reduces to the symmetric Gauss-Seidel step. M_SSOR is SPD
+    when A is SPD and ω ∈ (0, 2), so it is safe as a CG preconditioner.
+
+    SSOR outperforms scalar Jacobi on matrices with off-diagonal coupling —
+    FEM stiffness matrices, Poisson-like systems, and density-fitting metrics
+    with neighbour-basis coupling.
+
+    Args:
+        A: Symmetric positive definite matrix (n, n).
+        omega: Relaxation parameter in (0, 2). Default 1.0 (symmetric Gauss-Seidel).
+
+    Returns:
+        Callable M(r) → approx A^{-1} @ r for use as M= argument to `cg`.
+
+    Raises:
+        ValueError: If omega is not in (0, 2) or the diagonal has near-zero entries.
+    """
+    if not (0.0 < omega < 2.0):
+        raise ValueError(f"ssor_preconditioner: omega must be in (0, 2), got {omega}")
+    A, _ = _to_fp32(A)
+    d = torch.diagonal(A).clone()
+    if torch.any(d.abs() < 1e-15):
+        raise ValueError("ssor_preconditioner: diagonal has near-zero entries")
+    # Lower triangular factor: D + ω * L_strictly_lower
+    L_factor = omega * torch.tril(A, diagonal=-1) + torch.diag(d)
+    U_factor = L_factor.T  # valid for symmetric A
+    dscale = omega * (2.0 - omega) * d
+
+    def apply(r: torch.Tensor) -> torch.Tensor:
+        r_fp32, r_orig = _to_fp32(r)
+        t = torch.linalg.solve_triangular(L_factor, r_fp32.unsqueeze(-1), upper=False).squeeze(-1)
+        v = dscale * t
+        z = torch.linalg.solve_triangular(U_factor, v.unsqueeze(-1), upper=True).squeeze(-1)
+        return _restore(z, r_orig)
+
+    return apply
+
+
 def cg(
     A: torch.Tensor | Callable,
     b: torch.Tensor,