Skip to content

perf: in-place Cholesky buffer + copy-free numba solves for fnnls_cholesky - #453

Merged
Jammy2211 merged 2 commits into
mainfrom
feature/fnnls-inplace-cholesky-buffer
Aug 20, 2026
Merged

perf: in-place Cholesky buffer + copy-free numba solves for fnnls_cholesky#453
Jammy2211 merged 2 commits into
mainfrom
feature/fnnls-inplace-cholesky-buffer

Conversation

@Jammy2211

@Jammy2211 Jammy2211 commented Aug 20, 2026

Copy link
Copy Markdown
Collaborator

Closes #452.

What

The positive-only inversion solve (fnnls_cholesky) dominates the numba CPU Delaunay likelihood (~74% of a euclid eval on the production campaign fiducial — PyAutoLabs/autolens_profiling#151). The cost was implementation, not linear algebra: ~150 Bro–Jong active-set iterations, each rebuilding the ~1300² Cholesky factor via np.insert/np.delete (two O(n²) copies per change), plus scipy solve_triangular/cho_solve re-copying and finite-scanning the factor on every call.

Three changes, same mathematics:

  • In-place factor buffer — the factor lives in the top-left k×k corner of one preallocated buffer, grown by cholinsertlast_inplace and shrunk by choldeleteindexes_inplace; the existing numba Givens up/down-date kernels are unchanged (verified bitwise-stable on strided views). The out-of-place cholinsertlast/choldeleteindexes remain for reference and tests.
  • Copy-free numba solves_solve_upper_transposed_buffer (row-oriented forward substitution) and _cho_solve_buffer read the buffer's upper triangle directly, replacing the per-iteration scipy calls. This half is essential: with scipy still in the loop, its per-call copy of the strided view cancels the buffer's win (measured: no speedup with the buffer alone).
  • Numba shift kernel for deletes (_choldelete_shift_buffer) — numpy's overlapping slice assignments buffer the source region, re-introducing per-delete temp allocations; the kernel closes the deleted row/column gap with plain loops (bitwise-identical value movement).

Results (4-core cloud container, Delaunay Hilbert-1250 campaign fiducial)

Numerical equivalence

Results are not bitwise identical to the old path: LAPACK/our kernels pick layout-dependent but equally valid evaluation orders, and the old implementation already mixed memory layouts between its own iterations (scipy.linalg.cholesky returns F-order, np.insert/np.delete return C-order). Equivalence is enforced in the new test_autoarray/util/test_cholesky_inplace.py: cross-implementation agreement at 1e-13, exact factor-property checks (U′U = active submatrix), scipy.optimize.nnls cross-checks, and warm-start = cold-start. The degeneracy guard (_pivot_from_schur) and the producer-side non-finite check in fnnls_cholesky are unchanged.

Testing

  • test_autoarray/: 1026 passed, 51 skipped; the only failures are 3 pynufft transformer tests that fail identically on stock main in this environment (pre-existing).
  • New: test_autoarray/util/test_cholesky_inplace.py (18 tests).

Downstream

No public-API change — fnnls_cholesky(ZTZ, ZTx, P_initial) signature and caller (inversion_util.reconstruction_positive_only_from) are untouched. fix_constraint_cholesky's (module-internal) signature changed to thread the buffer. Likelihood values shift at the ~1e-9 relative level, within the established solver bistability tolerance (profiling pins use rtol 1e-6).

Block-pivoting / block-insertion phase-2 variants were prototyped and rejected (cycling / masked non-convergence on this system's near-degenerate columns) — findings in #452.

🤖 Generated with Claude Code

https://claude.ai/code/session_01Vcc7MUBMnNU6n8qqS9ioVZ

claude added 2 commits August 20, 2026 21:23
…lesky

The positive-only inversion solve (fnnls_cholesky) dominates the numba CPU
Delaunay likelihood (~74% of a euclid eval, autolens_profiling#151): ~150
Bro-Jong active-set iterations, each rebuilding the ~1300^2 factor via
np.insert/np.delete (two O(n^2) copies per change) and re-scanning/copying
it inside scipy's solve_triangular/cho_solve.

Two changes, same mathematics:

- The factor now lives in the top-left k x k corner of one preallocated
  buffer, grown/shrunk in place (cholinsertlast_inplace /
  choldeleteindexes_inplace; the numba Givens kernels are unchanged and
  proven bitwise-stable on strided views).
- The per-iteration triangular solves run through new copy-free numba
  kernels (_solve_upper_transposed_buffer, _cho_solve_buffer) that read the
  buffer's upper triangle directly — scipy on a strided view re-copies and
  finite-scans O(k^2) per call, which had cancelled the buffer's win.

On the real euclid 1310-param system (Delaunay Hilbert-1250 fiducial,
autolens_profiling PR #152): solve 3.29 s -> 1.40 s (2.35x), identical
active set, solutions agree to ~1e-9 absolute. Results are not bitwise
identical to the old path (LAPACK picks layout-dependent but equally valid
kernels; the old implementation already mixed layouts iteration-to-iteration)
— equivalence is tested at 1e-13 in the new
test_autoarray/util/test_cholesky_inplace.py, plus exact factor-property
checks and scipy.optimize.nnls cross-checks. The out-of-place
cholinsertlast/choldeleteindexes remain for reference and tests.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Vcc7MUBMnNU6n8qqS9ioVZ
Profiling the new solver showed choldeleteindexes_inplace as the largest
remaining cost (0.64 s of 1.45 s at euclid): numpy's overlapping slice
assignments buffer the source region, re-introducing per-delete temp
allocations. _choldelete_shift_buffer closes the deleted row/column gap
with plain numba loops (destinations strictly behind sources, upper
triangle only) — pure value movement, bitwise identical.

Euclid 1310-param system: solve 1.40 s -> 1.16 s (2.83x total vs the
np.insert/np.delete baseline's 3.29 s); euclid eval 2.34 s/call with the
profiling pin passing. test_autoarray: same 1026 passed / 3 pre-existing
pynufft failures.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01Vcc7MUBMnNU6n8qqS9ioVZ
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

perf: in-place Cholesky buffer + copy-free numba solves for fnnls_cholesky

2 participants