Repository navigation
Eliminate TrustRegion sparse per-step JᵀJ allocations - #1366
ChrisRackauckas-Claude wants to merge 4 commits into
Conversation
In the normal-form (and least-squares with normal-form damping) modes, `transpose(J) × J` on a SparseMatrixCSC went through MaybeInplace's sparse fallback, which materializes `copy(transpose(J))` and a fresh `spmatmul` product every iteration (~239 KB/step for LevenbergMarquardt with KLUFactorization on a 1000-dim tridiagonal problem). `Utils.normal_form_jacobian!!` now writes the values into the cached JᵀJ pattern using a transpose buffer and a dense accumulator held in the descent cache. The accumulation order matches SparseArrays' Gustavson spmatmul, so results are bitwise identical; a pattern mismatch falls back to the allocating product. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com> Agent-Harness: Claude Code Agent-Model: claude-opus-5-5 Agent-Session: https://claude.ai/code/session_01LPHREnnonfLg1VcE1EJovv
Reuse Utils.normal_form_jacobian!! / normal_form_workspace from SciML#1363 in MoreTrustRegionDescent so sparse TrustRegion no longer rebuilds JᵀJ via MaybeInplace's allocating CSC product each Jacobian update. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Cursor Agent <noreply@cursor.com> Agent-Harness: Cursor Agent CLI Agent-Model: auto Co-authored-by: Cursor <cursoragent@cursor.com>
48a1ff7 to
bd7e750
Compare
Independent review (Devin CLI 3000.11.3, model fusion-claude-opus-5-5-high-sidekick-swe-2-medium): CHANGES, risk low. Full reviewVERDICT: CHANGES Scope: the PR is stacked. Blocking findings
Non-blocking findings
What I ranAll on Julia 1.12.7 (aarch64). Two environments:
What I did not verify
Links
🤖 Posted by an AI agent — harness: Devin CLI 3000.11.3 (review), Claude Code 2.1.285 (posting) · model: fusion-claude-opus-5-5-high-sidekick-swe-2-medium |
Move `normal_form_workspace` and `normal_form_jacobian!!` from `NonlinearSolveBase.Utils` to the top-level module, declare them public next to the other normal-form traits, and document them in the developer docs. The SparseArrays extension now extends the public names, so the ExplicitImports `all_qualified_accesses_are_public` ignore entries for them are gone. Extend the TrustRegion sparse tests: `reinit!` with a new `p` on the default and KLU linsolves must match a fresh solve bitwise and keep `JᵀJ` sparse, and the warmed `step!` window must form a new Jacobian while allocating under 1000 bytes per step. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com> Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> Agent-Harness: Claude Code 2.1.301 Agent-Model: claude-opus-5-5 Agent-Session: local session ID 9ed6c56f-b474-4c6d-be68-63bf2ea757c3 (headless claude -p, no URL)
Summary
Stacked on #1363 (needs
Utils.normal_form_workspace/Utils.normal_form_jacobian!!inNonlinearSolveBaseSparseArraysExt).Moré
TrustRegionon sparse Jacobians rebuiltJᵀJevery Jacobian update through MaybeInplace's CSC__mul!(C .= A * B, ~239 KB/step on a 1000-dim tridiagonal Bratu problem). This PR builds the #1363 normal-form workspace once atMoreTrustRegionDescentinit and callsUtils.normal_form_jacobian!!on each update, so the Gramian stays in-place on a fixedjac_prototypeand falls back to the allocating product when an in-placejac!grows the pattern.Before / after (
JULIA_NUM_THREADS=1 taskset -c 120, 1000-dim sparse Bratu,TrustRegion()):step!max@allocated: 239536 → 144 bytes@benchmark solveminimum: 11.60 ms / 2993176 B → 11.94 ms / 1796536 B (time within machine noise; solve memory −1.2 MB)nsteps=38,njacs=5,nfactors=77,nsolve=114)239536 < 40000), passes with the fixjac!regression: unguarded fill leavesJᵀJ nnz=3/ wrong solve; guarded helper matches dense TR (maxerr=0)Test-group tails
GROUP=QANonlinearSolveBase:QA | 20 20 50.1s…Testing NonlinearSolveBase tests passedGROUP=QANonlinearSolveFirstOrder:QA | 20 20 56.5s…Testing NonlinearSolveFirstOrder tests passedGROUP=CoreNonlinearSolveBase:Testing NonlinearSolveBase tests passedGROUP=CoreNonlinearSolveFirstOrder:TrustRegion sparse normal-form | 8 8 12.6s…Testing NonlinearSolveFirstOrder tests passedNot verified: full
GROUP=All; GPU / non-CSC sparse layouts; banded/structured Jacobians on the Moré normal-form path; wall-time speedup beyond allocation (solve minima were noise-dominated on a loaded host).Reviewer pushback: whether stacking on #1363 is the right merge order; whether the residual ≤144 B/
@allocatednoise bound (< 40000) is the right regression threshold.Please ignore until reviewed by @ChrisRackauckas.
Risk assessment
Independent review: CHANGES addressed (reuse #1363 guarded helper; growing-pattern regression; Base QA)
🤖 Generated with Cursor Agent 2026.10.01-14929f9 (model: unknown (Cursor auto)), transcript /pr-tr-sparse-alloc/log.txt; orchestrated by Claude Code (claude-opus-5-5[1m]) https://claude.ai/code/session_01LPHREnnonfLg1VcE1EJovv
Made with Cursor
Review findings addressed
Commit
fcb0adcresponds to the Codex (gpt-6-astra) review and to the Devin review comment on this PR.Astra P2: the PR added new ExplicitImports
all_qualified_accesses_are_publicignore entries. Fixed.normal_form_workspaceandnormal_form_jacobian!!now live in the top-levelNonlinearSolveBasemodule instead ofUtils, and are declaredpublicnext torequires_normal_form_jacobian. Their docstrings state the extension contract, and they are listed indocs/src/devdocs/internal_interfaces.mdunder "Normal-Form Jacobian". The SparseArrays extension extends the public names, and both ignore entries are removed. I ran BaseGROUP=QAwith the ignore entries removed:QA | 19 1 20, withnormal_form_jacobian!! is not public in NonlinearSolveBase.Utilsand the same error fornormal_form_workspace;QA | 20 20.Devin: the new sparse path had no test of
reinit!with a changed parameter. Fixed.sparsity_tests__item4.jlnow runsinit → solve! → reinit!(cache, u0; p = p₂) → solve!withlinsolve = nothingand withKLUFactorization(). It asserts:< 1e-8;retcodeanduequal (==) to a freshsolve(remake(prob; p = p₂));JandJᵀJare both stillSparseMatrixCSC;JᵀJ == transpose(J) * J.This is new coverage, not a bug fix: these tests also pass with the
MoreTrustRegionDescentchange reverted.Devin: merge order. No code change. This branch still carries Form sparse JᵀJ in place in DampedNewtonDescent (LM normal-form step! allocations) #1363's commit (3cf0c9a), and this commit changes code that Form sparse JᵀJ in place in DampedNewtonDescent (LM normal-form step! allocations) #1363 introduced. Form sparse JᵀJ in place in DampedNewtonDescent (LM normal-form step! allocations) #1363 itself still has the two ignore entries. The options are: merge Form sparse JᵀJ in place in DampedNewtonDescent (LM normal-form step! allocations) #1363 first and then merge master into this branch (this commit then removes the entries on top of it), or merge this PR alone, which makes Form sparse JᵀJ in place in DampedNewtonDescent (LM normal-form step! allocations) #1363 redundant.
Devin (non-blocking): the allocation bound was not derived from anything. The bound is tightened from
< 40_000to< 1000, the same as the sibling LM test, and the test now also asserts thatcache.stats.njacsgrows inside the measured window.(144, 0, 0, 144, 0)with 2 new Jacobians on Julia 1.13.0, and(0, 0, 0, 0, 0)with 1 new Jacobian on 1.10.12.MoreTrustRegionDescentchange reverted to@bb cache.JᵀJ = transpose(J_) × J_, the test fails:Evaluated: 239536 < 1000.Devin (non-blocking): the top-of-file comment in item4 was retrospective. It now describes the property under test.
Devin (non-blocking): the TrustRegion Bratu solves return
Stalled/ShrinkThresholdExceededwith a residual of about 3e-11. Not changed. This is existing termination behaviour and also happens on the parent tree, so the new test asserts the residual rather than the retcode. It is worth a separate issue.Devin (non-blocking): the PR body's
nfactors/nsolvecounts differ from the reviewer's. Not changed. The counts depend on the Julia version, and the bitwise-identity claim was confirmed by the reviewer.Not breaking. The commit adds two public names to NonlinearSolveBase, so its next release needs a minor version bump rather than a patch. No existing public name changes. Following the repository's practice of separate release PRs, I did not bump the version.
Verification (for
fcb0adc)All of these ran locally with
JULIA_NUM_THREADS=4.Test groups
GROUP=CoreNonlinearSolveBase (Julia 1.13.0):Testing NonlinearSolveBase tests passed.GROUP=CoreNonlinearSolveFirstOrder (Julia 1.13.0):Sparse normal-form LevenbergMarquardt | 14 14,TrustRegion sparse normal-form | 21 21…Testing NonlinearSolveFirstOrder tests passed.GROUP=QANonlinearSolveBase:QA | 20 20on 1.13.0, andQA | 18 18on 1.10.12, where thepublicchecks do not run.Sparse test files on their own (item3 and item4)
14 14and21 21.MoreTrustRegionDescentchange reverted:TrustRegion sparse normal-form | 18 1 19, failing atEvaluated: 239536 < 1000.Docs and formatting
@docssection, in a minimal Documenter build withNonlinearSolveBaseandSparseArraysloaded. It had no errors, both docstrings rendered, and their@reflinks resolved. I did not run the fulldocs/make.jl, which needs Sundials, SciPy, PETSc and the MTK tutorials, plus linkcheck.--checkon the changed.jlfiles: exit 0.typoson the changed files: exit 0.Not verified: Julia
pre, the full docs build, GPU paths, and downstream packages.🤖 Posted by an AI agent — harness: Claude Code 2.1.301 · model: claude-opus-5-5
Conversation: local session ID 9ed6c56f-b474-4c6d-be68-63bf2ea757c3 (headless, no URL)
Risk assessment (independent review, codex-cli gpt-6.1-sol, reviewed head fcb0adc)
JULIA_DEPOT_PATH="$PWD/depot:~/.julia" JULIA_NUM_THREADS=4 JULIA_NUM_PRECOMPILE_TASKS=4 julia +release --project=env run-tests.jlon Julia 1.13.2: LM 14/14, TrustRegion 21/21, public extension access 3/3 passed. The driver includes the PR's item3/item4 tests in separate SafeTestsets, asserts both helpers are public, and runs ExplicitImports on the SparseArrays extension using exactly the existing repository qualified-access exclusions. This covers bitwise sparse products, rectangular matrices, changing patterns, allocation bounds, and reinit! with new parameters for default/KLU solvers. Initial and final CI checks both showed 18 pending, no failures: root https://github.com/SciML/NonlinearSolve.jl/actions/runs/38922327164 and sublibraries https://github.com/SciML/NonlinearSolve.jl/actions/runs/38922327037. Latest completed default-branch sublibrary CI passed: https://github.com/SciML/NonlinearSolve.jl/actions/runs/38812969499. Read default-branch root failure logs: both LTS jobs fail the existing PolyAlgorithm solve! inference test at core_tests__item3.jl:27, https://github.com/SciML/NonlinearSolve.jl/actions/runs/38812969394. No current-head failure exists to compare against it. Full Core/QA groups, LTS, docs build, downstreams, wall-time benchmarks, and an independent before-fix run were not performed in this review.git diff --checkpassed.🤖 Posted by an AI agent — harness: codex-cli · model: gpt-6.1-sol (independent reviewer; appended by Claude Code 2.1.301 · claude-opus-5-5)
Conversation: https://claude.ai/code/session_01VZjwf7Ux9YtTzwHGHMqnsq