SupernodalLU: sparse block-cache storage, AllocCheck/JET coverage, threaded race fix#1114
Conversation
|
Reworked per review to use LinearSolve's own default solver for the block caches, which is both simpler and strictly better than what I had. What changed. Each cached diagonal block is now Result — this is what my previous version could not do:
So the whole factorization type is inferable from its inputs, construction included. The JET test that was marked Memory, since holding all slots looks expensive: it isn't — the block buffer dominates. 35 796 B vs 34 264 B at np = 64, and 2 112 724 B vs 2 105 816 B at np = 512, i.e. 1.0×. One structural consequence:
Process note: while juggling branches I briefly mixed this work with an unrelated RFLU change in one working tree. The two are now split into disjoint commits (verified: TriangularSolve appears only in the other branch, zero times here), and both suites were re-run from a clean checkout of this branch — the results above are from that run, not the mixed one. |
|
Re: the But you were right that the Union was the wrong shape for a 99 %-empty vector. Pushed a sparse representation instead:
|
…readed race fix
Four changes to the vendored solver, found by auditing it against
PureKLU's optimization history and by writing the QA coverage.
1. Block caches are stored sparsely. Only blocks at or above
dense_threshold get a LinearSolve cache; everything narrower uses the
built-in kernel, which is the overwhelming majority - 0.6 % of
supernodes are cached on poisson2d_256 and 1.3 % on poisson3d_28,
where the median supernode is 5 and 1 columns wide. `bcaches` now
holds only the blocks that have one, with `bcacheidx[s]` indexing
into it (0 = built-in kernel), instead of a 99 %-empty full-length
vector with a nullability check on every lookup.
Caching every supernode instead was considered and rejected on the
numbers: a cache costs ~1476 B even for a 4x4 block, so poisson2d_256
alone would spend ~8 MB to make small blocks slower than the built-in
kernel - the same size effect behind PANEL_BLAS_CUTOFF.
The element type stays `Any`, and this is load-bearing rather than
laziness: a LinearCache type carries 13 parameters, and putting it in
SupernodalLUFactor's parameters pushes the sparse default solver's
cacheval - which holds a SupernodalLUFactor in one of its 30 slots -
past Julia's type-complexity limit, so inference widens it and
`@inferred init(prob, nothing)` fails for sparse matrices
(test/Core/default_algs.jl). Verified that a smaller concrete cache
type does not help either. Keeping the caches out of the type leaves
SupernodalLUFactor{Tv,Ti} fully concrete and inferable, at the cost
of one dispatch per cached supernode per factorization - tens of
calls against hundreds of ms of GEMM. The solve path never touches
the field.
2. A threaded data race, found by the new JET test. The per-worker
perturbation accumulator shared one variable with the serial phase,
so Julia boxed it and every spawned task incremented the same
counter; nperturbed could come out wrong, and it drives both the
refinement decision and the check=true singularity throw. The
existing threaded-equals-serial test missed it because it compares
factors and pivot sequences, and nperturbed was 0 on both sides.
3. The solve sweeps are now provably allocation-free. AllocCheck counts
the grow-on-demand resize! calls as possible allocations even when
they never fire, so the scratch size - derivable from the symbolic
analysis - is now recorded as SymbolicAnalysis.maxnu and applied once
at solve entry. The sweeps carry no growth branch at all, which also
removes a branch from the hot loop.
4. QA coverage. test/qa/allocations.jl gets a @check_allocs proof for
the sweeps (tridiagonal and 2D-mesh cases, the latter with maxnu=29
so panels are real) plus runtime zero-allocation assertions for
single- and multi-RHS solve!. test/qa/jet.jl asserts zero runtime
dispatch for solve! and _solve_panels! and that the factorization
type is concrete; factorization is marked broken with the reason,
since it carries the one deliberate dispatch described above.
Block factorization goes through LinearSolve's own default solver rather
than an algorithm resolved here, so it follows whatever the dense
default resolves to. src/SupernodalLU is included after src/default.jl
for that reason.
GROUP=Core and GROUP=QA both pass.
Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_017rr42T1vFRRT8D7DMAmJzf
950b68f to
e3af7ef
Compare
ChrisRackauckas-Claude
changed the title
2 participants
Please ignore until reviewed by @ChrisRackauckas.
Four changes to the vendored
SupernodalLUsolver, found by auditing it against PureKLU's optimization notes and by writing the QA coverage.1. Block caches are stored sparsely
Only blocks at or above
dense_thresholdget a LinearSolve cache; everything narrower uses the built-in kernel. That is nearly everything — measured on this branch:bcachesnow holds only the blocks that have one, withbcacheidx[s]indexing into it (0= built-in kernel), instead of a 99 %-empty full-length vector with a nullability check on every lookup.Caching every supernode was considered and rejected on the numbers: a cache costs 1476 B even for a 4×4 block (measured), so poisson2d_256 alone would spend ~8 MB to make small blocks slower than the built-in kernel — the same size effect that motivates
PANEL_BLAS_CUTOFF.The element type stays
Any, deliberately. ALinearCachecarries 13 type parameters, and putting it inSupernodalLUFactor's parameters pushes the sparse default solver's cacheval — which holds aSupernodalLUFactorin one of its 30 slots — past Julia's type-complexity limit. Inference then widens it and@inferred init(prob, nothing)fails for sparse matrices (test/Core/default_algs.jl), which is the property #1109 depends on. I tried a concrete narrow cache type too; it does not help, because the parameter count is the problem. Keeping caches out of the type leavesSupernodalLUFactor{Tv,Ti}fully concrete, at the cost of one runtime dispatch per cached supernode per factorization — tens of calls against hundreds of ms of GEMM. The solve path never touches the field.2. A threaded data race (found by the new JET test)
The per-worker perturbation accumulator shared one variable with the serial phase, so Julia boxed it and every spawned task incremented the same counter.
nperturbedcould come out wrong, and it drives both the refinement decision and thecheck=truesingularity throw. The existing threaded-equals-serial test missed it because it compares factors and pivot sequences, andnperturbedwas 0 on both sides.3. Solve sweeps are provably allocation-free
AllocCheck counts the grow-on-demand
resize!calls as possible allocations even when they never fire. The largest scratch size is derivable from the symbolic analysis, so it is now recorded asSymbolicAnalysis.maxnuand applied once at solve entry; the sweeps carry no growth branch at all, which also removes a branch from the hot loop.4. QA coverage
test/qa/allocations.jl: a@check_allocsproof for the sweeps, over a tridiagonal case (maxnu = 1) and a 2D mesh (maxnu = 29, so panels are genuinely exercised), plus runtime zero-allocation assertions for single- and multi-RHSsolve!.test/qa/jet.jl: zero-runtime-dispatch assertions forsolve!and_solve_panels!, and a concreteness assertion on the factorization type. Factorization (snlu/snlu!) is markedbroken = truewith the reason inline, since it carries the one deliberate dispatch described above — I did not weaken the solve-path assertions to accommodate it.Block factorization goes through LinearSolve's own default solver rather than an algorithm resolved here, so it follows whatever the dense default resolves to.
src/SupernodalLUis included aftersrc/default.jlfor that reason.GROUP=CoreandGROUP=QAboth pass locally.🤖 Generated with Claude Code
https://claude.ai/code/session_017rr42T1vFRRT8D7DMAmJzf