SupernodalLU: sparse block-cache storage, AllocCheck/JET coverage, threaded 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

Author: ChrisRackauckas

ext/LinearSolveSparseArraysExt.jl

@@ -704,28 +704,31 @@ function LinearSolve.init_cacheval(
     return nothing
 end
 
-const PREALLOCATED_SUPERNODAL = SNLU.snlu(
-    SparseMatrixCSC{Float64, Int64}(0, 0, [Int64(1)], Int64[], Float64[])
-)
-
 function LinearSolve.init_cacheval(
         alg::SupernodalLUFactorization, A::AbstractSparseArray{Float64, Int64}, b, u,
         Pl, Pr,
         maxiters::Int, abstol, reltol,
         verbose::Union{LinearVerbosity, Bool}, assumptions::OperatorAssumptions
     )
-    return PREALLOCATED_SUPERNODAL
+    return SNLU.snlu(
+        SparseMatrixCSC{Float64, Int64}(0, 0, [Int64(1)], Int64[], Float64[]);
+        dense_alg = alg.dense_alg
+    )
 end
 
-# SupernodalLU is pure Julia and factors any `Number` element type; the empty
-# cacheval carries the correct types so `pplu!`/`pplu` dispatch is concrete.
+# SupernodalLU is pure Julia and factors any `Number` element type.  The empty
+# prototype must resolve the same dense block algorithm as a real
+# factorization (hence `dense_alg = alg.dense_alg`), because the block-cache
+# type is part of the factorization type and the cacheval field is pinned to
+# whatever this returns.
 function LinearSolve.init_cacheval(
         alg::SupernodalLUFactorization, A::AbstractSparseArray{T, Ti}, b, u, Pl, Pr,
         maxiters::Int, abstol, reltol,
         verbose::Union{LinearVerbosity, Bool}, assumptions::OperatorAssumptions
     ) where {T <: Number, Ti <: Integer}
     return SNLU.snlu(
-        SparseMatrixCSC{T, Ti}(0, 0, [one(Ti)], Ti[], T[])
+        SparseMatrixCSC{T, Ti}(0, 0, [one(Ti)], Ti[], T[]);
+        dense_alg = alg.dense_alg
     )
 end
 

src/LinearSolve.jl

@@ -435,7 +435,6 @@ function defaultalg_symbol end
 
 include("verbosity.jl")
 include("blas_logging.jl")
-include("SupernodalLU/SupernodalLU.jl")
 include("generic_lufact.jl")
 include("eigenvalue.jl")
 include("common.jl")
@@ -453,6 +452,9 @@ include("preconditioners.jl")
 include("preferences.jl")
 include("solve_function.jl")
 include("default.jl")
+# after default.jl: the vendored solver caches its dense diagonal blocks
+# with LinearSolve's own default solver, so it needs DefaultLinearSolver{,Init}
+include("SupernodalLU/SupernodalLU.jl")
 include("init.jl")
 include("adjoint.jl") # LinearSolveAdjoint struct definition only; rrules are in ChainRulesCore ext
 

src/SupernodalLU/interface.jl

@@ -110,36 +110,40 @@ end
 # RFLU when RecursiveFactorization is loaded, MKL/LAPACK/generic otherwise;
 # resolved once at analysis time).  Small blocks and generic element types
 # use the built-in static-perturbation kernel.
+# One dense block cache over an `np x np` buffer.  `dense_alg === nothing`
+# means LinearSolve's default solver, which is the point: its cache type is
+# `LinearCache{..., DefaultLinearSolver, DefaultLinearSolverInit{...}}`, which
+# holds a slot for every candidate algorithm and picks between them with a
+# runtime enum.  The type therefore depends only on the element type - not on
+# which algorithm the default happens to select on this machine - so the whole
+# factorization type is inferable from its inputs and `init_cacheval`'s empty
+# prototype matches every real factorization by construction.  The extra slots
+# cost nothing measurable: the block buffer dominates (35 796 B vs 34 264 B at
+# np = 64, 2 112 724 B vs 2 105 816 B at np = 512).
+function _block_cache(::Type{Tv}, alg, np::Int) where {Tv}
+    return SciMLBase.init(
+        LinearSolve.LinearProblem(Matrix{Tv}(undef, np, np), zeros(Tv, np)),
+        alg;
+        alias = SciMLBase.LinearAliasSpecifier(alias_A = true, alias_b = true),
+        assumptions = LinearSolve.OperatorAssumptions(true)
+    )
+end
+
+
 function _dense_block_caches(
-        W::Vector{Matrix{Tv}}, sym::SymbolicAnalysis, dense_alg, dense_threshold::Int
+        ::Type{Tv}, sym::SymbolicAnalysis, alg, dense_threshold::Int
     ) where {Tv}
     nsuper = length(sym.sstart) - 1
-    bcaches = Vector{Any}(nothing, nsuper)
-    Tv <: LinearAlgebra.BlasFloat || return bcaches
-    assump = LinearSolve.OperatorAssumptions(true)
+    bcaches = Any[]
+    bcacheidx = zeros(Int, nsuper)
+    Tv <: LinearAlgebra.BlasFloat || return bcaches, bcacheidx
     for s in 1:nsuper
         np = sym.sstart[s + 1] - sym.sstart[s]
         np >= dense_threshold || continue
-        B = Matrix{Tv}(undef, np, np)   # cache-owned block buffer
-        b = zeros(Tv, np)
-        alg = dense_alg
-        if alg === nothing
-            # `A === nothing` is LinearSolve's documented stand-in for "dense
-            # matrix" in `defaultalg` (our block is a strided view, which the
-            # dense branch does not match on), so this resolves exactly the
-            # dense default for this block's size and element type.
-            da = LinearSolve.defaultalg(nothing, b, assump)
-            alg = da isa LinearSolve.DefaultLinearSolver ?
-                LinearSolve.algchoice_to_alg(Symbol(da.alg)) : da
-        end
-        bcaches[s] = SciMLBase.init(
-            LinearSolve.LinearProblem(B, b),
-            alg;
-            alias = SciMLBase.LinearAliasSpecifier(alias_A = true, alias_b = true),
-            assumptions = assump
-        )
+        push!(bcaches, _block_cache(Tv, alg, np))
+        bcacheidx[s] = length(bcaches)
     end
-    return bcaches
+    return bcaches, bcacheidx
 end
 
 # Allocate all numeric state (panels, V/Vt + position maps, every workspace)
@@ -166,6 +170,22 @@ function _snlu_build(
     end
     V = M[sym.qf, sym.qf]
     Vt, Tpos = _transpose_map(V)
+    return _snlu_assemble(
+        sym, A, M, ms, W, Z, rowsfac, V, Vt, Tpos, dense_alg, dense_threshold,
+        eps_pivot, threaded, check
+    )
+end
+
+function _snlu_assemble(
+        sym::SymbolicAnalysis, A::SparseMatrixCSC{Tv, Ti},
+        M::SparseMatrixCSC{Tv, Ti}, ms::Union{MatchScale, Nothing},
+        W::Vector{Matrix{Tv}}, Z::Vector{Matrix{Tv}}, rowsfac::Vector{Vector{Int}},
+        V::SparseMatrixCSC{Tv, Ti}, Vt::SparseMatrixCSC{Tv, Ti}, Tpos::Vector{Int},
+        dense_alg, dense_threshold::Int, eps_pivot::Float64, threaded::Bool,
+        check::Bool
+    ) where {Tv, Ti <: Integer}
+    n = sym.n
+    bcaches, bcacheidx = _dense_block_caches(Tv, sym, dense_alg, dense_threshold)
     F = SupernodalLUFactor{Tv, Ti}(
         sym, W, Z, collect(1:n), collect(1:n), rowsfac, 0, NaN, eps_pivot,
         A, Vector{Tv}(undef, n), Vector{Tv}(undef, 64),
@@ -178,7 +198,7 @@ function _snlu_build(
         Matrix{Tv}(undef, n, 0),
         Vector{Tv}(undef, n), Vector{Tv}(undef, n), Vector{Tv}(undef, n),
         threaded,
-        _dense_block_caches(W, sym, dense_alg, dense_threshold)
+        bcaches, bcacheidx
     )
     _build_vpos!(F, A)
     _factor_core!(F)

src/SupernodalLU/numeric.jl

@@ -77,16 +77,29 @@ mutable struct SupernodalLUFactor{Tv, Ti <: Integer}
     ir_dx::Vector{Tv}                # iterative-refinement correction
     btmp::Vector{Tv}                 # in-place ldiv! RHS copy
     threaded::Bool                   # use the etree-parallel numeric phase
-    # Dense LinearSolve caches for large diagonal blocks (`nothing` = built-in
-    # kernel).  Deliberately `Vector{Any}`: the solve path never touches it and
-    # `_cache_lu!` is a function barrier, so the only dynamic dispatch is one
-    # call per cached supernode per factorization (tens of calls against
-    # hundreds of ms of GEMM).  Narrowing it to `Vector{Union{Nothing, C}}` for
-    # union-splitting was tried and measured 2.3-2.7 % *slower* on
-    # poisson3d_28/poisson2d_256, and it breaks `init_cacheval`'s cacheval
-    # type-pinning (the empty prototype has no caches, so its element type
-    # cannot match a real factorization's).  Do not retry.
+    # Dense LinearSolve caches for the diagonal blocks at or above
+    # `dense_threshold`; everything narrower uses the built-in kernel, which is
+    # the overwhelming majority (0.6 % of supernodes are cached on
+    # poisson2d_256, 1.3 % on poisson3d_28 - the median supernode is 1-5
+    # columns wide).  Stored sparsely: `bcaches` holds only the blocks that
+    # have one and `bcacheidx[s]` indexes into it, 0 meaning "no cache", so
+    # there is no mostly-empty full-length vector and no nullability check on
+    # the lookup.
+    #
+    # The element type is deliberately `Any`.  Each entry is a `LinearCache`
+    # over the block, and a `LinearCache` type carries 13 parameters; putting
+    # it in `SupernodalLUFactor`'s own 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).  Keeping the caches out of the type leaves
+    # `SupernodalLUFactor{Tv, Ti}` fully concrete and inferable from its
+    # inputs, which matters far more: the solve path never touches this field,
+    # and `_cache_lu!` is a function barrier, so the only dynamic dispatch is
+    # one call per cached supernode per factorization - tens of calls against
+    # hundreds of ms of GEMM.
     bcaches::Vector{Any}
+    bcacheidx::Vector{Int}
 end
 
 nperturbed(F::SupernodalLUFactor) = F.nperturbed
@@ -237,9 +250,9 @@ function _dense_block_factor!(
         F, s::Int, np::Int, epsv, ipiv::AbstractVector{Int}, scratch
     )
     Ws = F.W[s]
-    bc = F.bcaches[s]
-    bc === nothing && return _block_lu!(Ws, np, epsv, ipiv)
-    ok = _cache_lu!(bc, Ws, np, epsv, ipiv)
+    i = F.bcacheidx[s]
+    i == 0 && return _block_lu!(Ws, np, epsv, ipiv)
+    ok = _cache_lu!(F.bcaches[i], Ws, np, epsv, ipiv)
     ok && return 0
     @inbounds for j in 1:np
         ipiv[j] = j
@@ -259,7 +272,7 @@ function _cache_lu!(
     end
     bc.isfresh = true
     SciMLBase.solve!(bc)
-    fact = _lu_from_cacheval(bc.cacheval)
+    fact = _block_lu_of(bc)
     fac = fact.factors
     @inbounds for j in 1:np
         abs(fac[j, j]) >= epsv || return false
@@ -294,6 +307,35 @@ end
 _lu_from_cacheval(cv::LinearAlgebra.LU) = cv
 _lu_from_cacheval(cv::Tuple) = _lu_from_cacheval(first(cv))
 
+# The default solver keeps one slot per candidate algorithm and records the
+# choice in a runtime enum, so pull the factorization out of the slot that was
+# actually used.  Written as a literal-symbol chain rather than
+# `getfield(cv, Symbol(choice))` so every branch stays concrete; the result is
+# a small union of `LU` types, which Julia union-splits.
+function _lu_from_cacheval(cv::LinearSolve.DefaultLinearSolverInit, choice)
+    DAC = LinearSolve.DefaultAlgorithmChoice
+    if choice === DAC.LUFactorization
+        return _lu_from_cacheval(cv.LUFactorization)
+    elseif choice === DAC.RFLUFactorization
+        return _lu_from_cacheval(cv.RFLUFactorization)
+    elseif choice === DAC.MKLLUFactorization
+        return _lu_from_cacheval(cv.MKLLUFactorization)
+    elseif choice === DAC.AppleAccelerateLUFactorization
+        return _lu_from_cacheval(cv.AppleAccelerateLUFactorization)
+    elseif choice === DAC.GenericLUFactorization
+        return _lu_from_cacheval(cv.GenericLUFactorization)
+    elseif choice === DAC.BLISLUFactorization
+        return _lu_from_cacheval(cv.BLISLUFactorization)
+    end
+    throw(ArgumentError("SupernodalLU: dense block solver resolved to $choice, which does not expose an LU factorization; pass `dense_alg` explicitly"))
+end
+
+# Factorization behind a block cache, whichever solver the cache is using.
+@inline _block_lu_of(bc) = _lu_from_cacheval(bc.cacheval)
+@inline function _block_lu_of(bc::LinearSolve.LinearCache{<:Any, <:Any, <:Any, <:Any, <:LinearSolve.DefaultLinearSolver})
+    return _lu_from_cacheval(bc.cacheval, bc.alg.alg)
+end
+
 # Transpose of V with a position map: Vt.nzval[q] == V.nzval[Tpos[q]] forever
 # (structure is static), so refactorization refreshes Vt by one gather pass.
 function _transpose_map(V::SparseMatrixCSC{Tv, Ti}) where {Tv, Ti}
@@ -677,13 +719,16 @@ function _core_threaded!(F::SupernodalLUFactor{Tv}, epsv) where {Tv}
                     relmap = zeros(Int, sym.n)
                     ipiv = Vector{Int}(undef, sym.n)
                     buf = Vector{Tv}(undef, 64)
-                    acc = 0
+                    # `local`: without it this shares the enclosing function's
+                    # `acc`, which boxes it (killing inference) and races every
+                    # worker against the others' perturbation counts.
+                    local nloc = 0
                     for (lo, hi) in q
                         for s in lo:hi
-                            acc += _process_supernode!(F, s, epsv, relmap, ipiv, buf)
+                            nloc += _process_supernode!(F, s, epsv, relmap, ipiv, buf)
                         end
                     end
-                    Threads.atomic_add!(npert_par, acc)
+                    Threads.atomic_add!(npert_par, nloc)
                 end
             end
         end

src/SupernodalLU/solve.jl

@@ -85,7 +85,17 @@ end
     return nothing
 end
 
+# Ensure the panel scratch can hold the widest update-row set times `nrhs`.
+# Called once per solve so the sweeps below carry no growth branch and are
+# provably allocation-free (see the AllocCheck test in test/qa/allocations.jl).
+@inline function _ensure_panel_scratch!(F::SupernodalLUFactor, nrhs::Int)
+    need = F.sym.maxnu * nrhs
+    length(F.gbuf) < need && resize!(F.gbuf, need)
+    return nothing
+end
+
 # y := U \ (L \ (P * y)) in permuted space; y enters as V-row-ordered rhs.
+# Precondition: `_ensure_panel_scratch!(F, 1)` has been called.
 function _solve_panels!(y::AbstractVector{Tv}, F::SupernodalLUFactor{Tv}) where {Tv}
     sym = F.sym
     sstart = sym.sstart
@@ -105,7 +115,6 @@ function _solve_panels!(y::AbstractVector{Tv}, F::SupernodalLUFactor{Tv}) where
             ldiv!(UnitLowerTriangular(view(Ws, 1:np, 1:np)), xb)
         end
         if nu > 0
-            length(buf) < nu && resize!(buf, max(nu, 2 * length(buf)))
             t = view(buf, 1:nu)
             if np < PANEL_BLAS_CUTOFF
                 _panel_gemv!(t, Ws, xb, np, nu)
@@ -126,7 +135,6 @@ function _solve_panels!(y::AbstractVector{Tv}, F::SupernodalLUFactor{Tv}) where
         Ws = F.W[s]
         xb = view(y, c1:c2)
         if nu > 0
-            length(buf) < nu && resize!(buf, max(nu, 2 * length(buf)))
             t = view(buf, 1:nu)
             @simd for k in 1:nu
                 t[k] = y[R[k]]
@@ -147,6 +155,7 @@ function _solve_panels!(y::AbstractVector{Tv}, F::SupernodalLUFactor{Tv}) where
 end
 
 # Multi-RHS variant: same sweeps with gemm-shaped updates on an n×nrhs block.
+# Precondition: `_ensure_panel_scratch!(F, nrhs)` has been called.
 function _solve_panels!(Y::AbstractMatrix{Tv}, F::SupernodalLUFactor{Tv}) where {Tv}
     sym = F.sym
     sstart = sym.sstart
@@ -163,7 +172,6 @@ function _solve_panels!(Y::AbstractMatrix{Tv}, F::SupernodalLUFactor{Tv}) where
         Yb = view(Y, c1:c2, :)
         ldiv!(UnitLowerTriangular(view(Ws, 1:np, 1:np)), Yb)
         if nu > 0
-            length(buf) < nu * nrhs && resize!(buf, max(nu * nrhs, 2 * length(buf)))
             T = reshape(view(buf, 1:(nu * nrhs)), nu, nrhs)
             mul!(T, view(Ws, (np + 1):(np + nu), 1:np), Yb)
             for r in 1:nrhs, k in 1:nu
@@ -180,7 +188,6 @@ function _solve_panels!(Y::AbstractMatrix{Tv}, F::SupernodalLUFactor{Tv}) where
         Ws = F.W[s]
         Yb = view(Y, c1:c2, :)
         if nu > 0
-            length(buf) < nu * nrhs && resize!(buf, max(nu * nrhs, 2 * length(buf)))
             T = reshape(view(buf, 1:(nu * nrhs)), nu, nrhs)
             for r in 1:nrhs, k in 1:nu
                 T[k, r] = Y[R[k], r]
@@ -203,6 +210,7 @@ function _solve_once!(x::AbstractVector{Tv}, F::SupernodalLUFactor{Tv}, b::Abstr
     rp = F.rowperm
     Rs = F.Rs
     Cs = F.Cs
+    _ensure_panel_scratch!(F, 1)
     @inbounds for k in 1:n
         i = rp[p[k]]
         y[k] = Rs[i] * b[i]
@@ -229,6 +237,7 @@ function _solve_once!(X::AbstractMatrix{Tv}, F::SupernodalLUFactor{Tv}, B::Abstr
     n = F.sym.n
     nrhs = size(B, 2)
     Y = _scratch_mat!(F, nrhs)
+    _ensure_panel_scratch!(F, nrhs)
     p = F.p
     qf = F.sym.qf
     rp = F.rowperm

src/SupernodalLU/symbolic.jl

@@ -185,6 +185,7 @@ struct SymbolicAnalysis
     rows::Vector{Vector{Int}}  # update rows (> last column) of each supernode
     snof::Vector{Int}        # column -> supernode
     nnzL::Int                # entries in L panels incl. unit diagonal & padding
+    maxnu::Int               # widest update-row set over all supernodes
     ccp::Vector{Int}         # contribution schedule: target s gets ccp[s]:ccp[s+1]-1
     cd::Vector{Int}          # ... from descendant cd[k]
     cj1::Vector{Int}         # ... starting at rows[cd[k]][cj1[k]]
@@ -396,5 +397,6 @@ function snlu_symbolic(
         end
     end
     ccp, cd, cj1 = _contrib_schedule(sstart, rows, snof)
-    return SymbolicAnalysis(n, qf, parentF, sstart, rows, snof, nnzL, ccp, cd, cj1)
+    maxnu = ns == 0 ? 0 : maximum(length, rows)
+    return SymbolicAnalysis(n, qf, parentF, sstart, rows, snof, nnzL, maxnu, ccp, cd, cj1)
 end