Skip to content

trnsparse 0.3.2 — CG + power iteration on BSR

Choose a tag to compare

@scttfrdmn scttfrdmn released this 14 Apr 23:18
· 25 commits to main since this release

Phase 1 plumbing of #22 on-chip iterative solvers. Python-level CG and power iteration on top of `bsr_spmm`.

Added

  • `cg_bsr(A, b, x0, tol, max_iter, M=None) -> (x, iters, rel)` — Conjugate Gradient for SPD BSR matrices. Takes an optional preconditioner.
  • `power_iteration_bsr(A, v0, max_iter, tol) -> (lam, v, iters)` — dominant eigenpair via power iteration.
  • `jacobi_preconditioner_bsr(A)` — diagonal preconditioner builder.
  • `bsr_diagonal(A)` — main-diagonal extractor.
  • `docs/iterative_solvers.md` — design note covering the v0.3.2 plumbing and the v0.4.0 fused-kernel goal.
  • `tests/test_iterative.py` (8 tests, scipy parity at `atol=1e-4`).
  • `benchmarks/bench_iterative.py` — at 128×128 SPD: scipy 310 μs, trnsparse 369 μs (1.19×).

Not in this release

The architectural win from #22's acceptance list — A SBUF-resident across all iterations — requires a fused NKI kernel that wraps the CG loop. That's tracked in #24 for v0.4.0. Today each CG iteration dispatches one `bsr_spmm`, so A round-trips to HBM per iter on NKI. The API stays stable across the v0.4.0 transition — users get the fused-kernel speedup automatically.

Validation

  • 45/45 CPU tests pass (37 existing + 8 new iterative).
  • Simulator CI (`nki-simulator`) green on ubuntu-latest.
  • No NKI code changes — `bsr_spmm` unchanged; hardware re-validation not required.

Closes #22 Phase 1. v0.4.0 tracker: #24.