Skip to content

fix: don't rebuild the linear solver on every inline linear SCC solve - #160

Merged
ChrisRackauckas merged 5 commits into
JuliaComputing:mainfrom
AJ0070:fix/inline-linear-scc-linsolve-cache
Sep 16, 2026
Merged

ChrisRackauckas merged 5 commits into
JuliaComputing:mainfrom
AJ0070:fix/inline-linear-scc-linsolve-cache

Conversation

@AJ0070

@AJ0070 AJ0070 commented Sep 11, 2026

Copy link
Copy Markdown
Contributor

Closes #157.

safe_ldiv was building a fresh LinearProblem and LinearCache on every call, which is 23 allocations per solve no matter how small the system is. It now keeps the LinearCache in task local storage and reuses it, so a steady state call neither builds a LinearProblem nor allocates a factorization.

One cache per eltype and size. Two solve sites of the same shape can share one, so the solution gets copied back into b, which belongs to that site alone, rather than handing back the cache's own buffer. The lookup is necessarily type unstable so the solve sits behind a function barrier.

Task local rather than kept alongside the system because the emitted expression is shared by every problem built from it, while A and b are not. With one shared cache and 8 threads on their own problems, 18289 of 24000 results come back wrong; task local gives 0.

The new values are filled into the cache's own buffers and then assigned back rather than written through cache.A in place, so that LinearSolve runs its invalidation. Writing in place keeps the partials of the first solve and silently returns stale derivatives under ForwardDiff on every later call, with the values still correct. There's a test for it now.

Allocations are flat at 112 B / 3 per RHS call regardless of size, and it's faster than main throughout: 144 vs 173 ns at N=2, 392 vs 557 at N=8, 4438 vs 5186 at N=32. The cached path still picks the same factorization as the uncached one, GenericLU at m=4 and AppleAccelerateLU at m=32.

@AJ0070
AJ0070 force-pushed the fix/inline-linear-scc-linsolve-cache branch from ee69471 to f6accf0 Compare September 11, 2026 14:33

@DhairyaLGandhi DhairyaLGandhi left a comment

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Would be good to see how the allocations are looking like now for a representative system. Couple thoughts:

Comment thread lib/ModelingToolkitTearing/src/reassemble.jl Outdated
Comment thread lib/ModelingToolkitTearing/src/reassemble.jl Outdated
@AJ0070

AJ0070 commented Sep 11, 2026

Copy link
Copy Markdown
Contributor Author

You're right. I swept SCC sizes and going straight to lu! only wins for tiny systems:

 N   main                       this PR
 2   1424 B / 25 al   186 ns     272 B / 6 al   161 ns   1.16x
 8   2608 B / 25 al   469 ns     800 B / 6 al   640 ns   0.73x
16   5968 B / 27 al  1502 ns    2400 B / 7 al  2740 ns   0.55x
32  18832 B / 27 al  5256 ns    8672 B / 7 al  7424 ns   0.71x

It gives up the default's backend switch (GenericLU below ~8, AppleAccelerateLU above),
so it's nearly 2x slower at N=16. Reworking to init + solve!.

One steer I need: where should the cache live? In the expression it ends up shared across
every problem built from the system, which the A/b diffcaches aren't today.
MTKParameters is the right scope but add_diffcache only hands back a numeric buffer, so
it needs a new slot type in MTKBase. Happy to add that, or use a task-local cache keyed per
solve site if you'd rather keep it here.

Two things I hit prototyping: cache.A = A doesn't work on the AD path since
DualLinearCache wants a concrete Matrix, so it has to be copyto!, and we need one
cache per element type.

@DhairyaLGandhi

Copy link
Copy Markdown

I think keeping it task local would be cleaner. Do you have an example sketch of what you're running into with the factorization getting shared?

Ideally we would want to re-use as much as we can, while remaining correct and fast.

@AJ0070

AJ0070 commented Sep 11, 2026

Copy link
Copy Markdown
Contributor Author

Sketch: with one cache shared by everything and 8 threads each running its own problem from the same system, 18289 of 24000 results come back wrong. Task local gives 0, same as main today. The A and b buffers are already per problem through the diffcaches, so it's only the factorization that would be shared.

Gone with task local then. One cache per eltype and size, and since two solve sites of the same shape can share one, I copy the solution back into b rather than handing back the cache's own buffer. The lookup has to be type unstable so the solve sits behind a function barrier, which mattered more than I expected, without it small systems came out slower than main.

Allocations are flat at 112 B / 3 regardless of size now, and it's faster than main throughout: 109 vs 173 ns at N=2, 373 vs 557 at N=8, 4180 vs 5186 at N=32. Checked that the cached path still picks the same factorization as the uncached one, GenericLU at m=4 and AppleAccelerateLU at m=32.

Comment thread lib/ModelingToolkitTearing/src/reassemble.jl Outdated
Comment thread lib/ModelingToolkitTearing/src/reassemble.jl Outdated
Comment thread lib/ModelingToolkitTearing/src/reassemble.jl Outdated
Comment thread lib/ModelingToolkitTearing/src/reassemble.jl Outdated
Comment thread lib/ModelingToolkitTearing/src/reassemble.jl Outdated
@AJ0070
AJ0070 marked this pull request as draft September 14, 2026 17:44
@ChrisRackauckas
ChrisRackauckas merged commit ccca981 into JuliaComputing:main Sep 16, 2026
9 checks passed
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.

safe_ldiv builds a fresh LinearCache per call: 2.6 kB and 39 allocations per RHS evaluation of an inline-linear-SCC model, 1 MB per 10 ms control tick

3 participants