Skip to content

SupernodalLU: concretely typed block caches + AllocCheck/JET coverage (fixes a threaded data race)#1114

Draft
ChrisRackauckas-Claude wants to merge 1 commit into
SciML:mainfrom
ChrisRackauckas-Claude:supernodal-lu-inference
Draft

SupernodalLU: concretely typed block caches + AllocCheck/JET coverage (fixes a threaded data race)#1114
ChrisRackauckas-Claude wants to merge 1 commit into
SciML:mainfrom
ChrisRackauckas-Claude:supernodal-lu-inference

Conversation

@ChrisRackauckas-Claude

Copy link
Copy Markdown
Member

⚠️ Draft — please ignore until reviewed by @ChrisRackauckas.

Independent of #1109, #1111, #1112 and #1113.

What

Replaces the runtime-typed bcaches::Vector{Any} with a concrete Vector{Union{Nothing,C}} (Vector{Nothing} for element types that never cache a block), and adds the AllocCheck/JET coverage that change deserves.

How the type gets pinned down. The block-cache type is now a function of (Tv, Ti, dense_alg) alone: the dense algorithm is resolved once from the element type instead of per block width. LinearCache carries no size parameter, so one algorithm gives one cache type for every block, and init_cacheval can build a prototype whose type matches every real factorization — the same contract PureUMFPACKFactorization follows, and precisely what made a naive narrowing fail on an earlier attempt (the empty prototype has no caches, so its element type could never match). init_cacheval now also forwards alg.dense_alg, so an explicitly chosen dense algorithm matches too.

Construction resolves that algorithm at runtime by design — it follows LinearSolve's dense default, which depends on what is loaded — so the choice is funnelled through one function barrier. Scoped to the solver module, JET reports:

path runtime dispatches
solve! 0
snlu! (refactorization) 0
_solve_panels! 0
snlu (one-time construction) 3 (was 8)

Two real defects fell out of writing the tests

JET found a data race. The threaded factorization's per-worker accumulator shared one variable with the serial phase. Julia boxed it (what JET flagged) and every spawned task incremented the same counter — so nperturbed could come out wrong, and it feeds both the refinement decision and the check=true singularity throw. The existing threaded≡serial test missed it because it compares factors and pivot sequences, and nperturbed was 0 on both sides. Fixed with a local on a distinctly named accumulator.

AllocCheck forced a design fix. It could not prove the sweeps allocation-free: the grow-on-demand resize! calls count as possible allocations even when they never fire. The largest scratch size is derivable from the symbolic analysis, so SymbolicAnalysis now carries maxnu, sizing happens once at solve entry, and the sweeps carry no growth branch at all. @check_allocs now proves _solve_panels! allocation-free, and the hot loops lost a branch.

Tests

  • test/qa/allocations.jl: 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: @test_opt for solve!/snlu!/_solve_panels! — scoped with target_modules, as elsewhere in that file, so Base error-path noise (show/AssertionError reachable from sparse indexing) does not mask our own code — concreteness assertions on the factorization type, and construction marked broken with the reason stated.

Measured perf-neutral: refactorization 0.1676 s vs 0.1703 s (poisson3d_28), 0.0883 s vs 0.0903 s (poisson2d_256).

GROUP=Core and GROUP=QA both pass.

🤖 Generated with Claude Code

https://claude.ai/code/session_017rr42T1vFRRT8D7DMAmJzf

…coverage

Replaces the runtime-typed bcaches::Vector{Any} with a concrete
Vector{Union{Nothing,C}} (Vector{Nothing} for element types that never
cache a block), and adds the QA coverage that the change deserves.

How the type is pinned down. The block-cache type is now a function of
(Tv, Ti, dense_alg) alone: the dense algorithm is resolved once from the
element type instead of per block width. LinearCache carries no size
parameter, so one algorithm gives one cache type for every block, and
init_cacheval can build a prototype whose type matches every real
factorization - the same contract PureUMFPACKFactorization follows, and
the thing that made a naive narrowing fail before (the empty prototype
has no caches, so its element type could never match). init_cacheval now
also forwards alg.dense_alg so an explicitly chosen dense algorithm
matches too.

Construction resolves that algorithm at runtime by design (it follows
LinearSolve's dense default, which depends on what is loaded), so the
choice is funnelled through one function barrier: the factorization body
is fully typed and the resulting SupernodalLUFactor{Tv,Ti,C} is concrete.
Scoped to the solver module, JET reports 0 runtime dispatches for solve!,
snlu! and _solve_panels! - the repeated-use paths - and construction drops
from 8 to 3.

Two real defects fell out of writing the tests:

- JET found that the threaded factorization's per-worker accumulator
  shared one variable with the serial phase. Julia boxed it (which is
  what JET flagged) AND every spawned task incremented the same counter,
  so nperturbed could come out wrong - it feeds 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. Fixed with a
  `local` on a distinctly named accumulator.

- AllocCheck could not prove the solve sweeps allocation-free: the
  grow-on-demand resize! calls count as possible allocations even when
  they never fire. The largest scratch size is derivable from the
  symbolic analysis, so SymbolicAnalysis now carries maxnu, sizing
  happens once at solve entry, and the sweeps carry no growth branch at
  all. @check_allocs now proves _solve_panels! allocation-free, and the
  hot loops lost a branch.

Tests: test/qa/allocations.jl gets a @check_allocs proof for the sweeps
plus runtime zero-allocation assertions for single- and multi-RHS
solve!; test/qa/jet.jl gets @test_opt (scoped to the solver module, as
elsewhere in that file, so Base error-path noise does not mask our code)
for solve!/snlu!/_solve_panels!, concreteness assertions on the
factorization type, and construction marked broken with the reason.

Measured perf-neutral: refactorization 0.1676s vs 0.1703s (poisson3d_28)
and 0.0883s vs 0.0903s (poisson2d_256).

GROUP=Core and GROUP=QA both pass.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
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.

2 participants