SupernodalLU: concretely typed block caches, plus AllocCheck and JET 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>

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/SupernodalLU/interface.jl

@@ -110,34 +110,55 @@ 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.
+# Size at which the dense default is probed.  The dense algorithm is resolved
+# from the element type alone rather than per block, so the block-cache type -
+# and hence the factorization type - is a function of `(Tv, Ti, dense_alg)`
+# only.  That is what lets `init_cacheval` build a prototype whose type matches
+# every real factorization (the same contract `PureUMFPACKFactorization`
+# follows), and it keeps `bcaches` concretely typed.  `LinearCache` carries no
+# size parameter, so one algorithm gives one cache type for every block width.
+const DENSE_ALG_PROBE_SIZE = 128
+
+# The dense algorithm used for every cached diagonal block.  `A === nothing` is
+# LinearSolve's documented stand-in for "dense matrix" in `defaultalg` (a panel
+# block is a strided buffer, and we want the dense branch), so this resolves
+# exactly what the dense default would pick for this element type.
+function _resolve_dense_alg(::Type{Tv}, dense_alg) where {Tv}
+    Tv <: LinearAlgebra.BlasFloat || return nothing
+    dense_alg === nothing || return dense_alg
+    b = zeros(Tv, DENSE_ALG_PROBE_SIZE)
+    da = LinearSolve.defaultalg(nothing, b, LinearSolve.OperatorAssumptions(true))
+    return da isa LinearSolve.DefaultLinearSolver ?
+        LinearSolve.algchoice_to_alg(Symbol(da.alg)) : da
+end
+
+# One dense block cache over an `np x np` buffer.
+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
+
+# Concrete element type of `bcaches` for this element type and algorithm.
+# `Nothing` when no block will ever be cached (non-BLAS element types).
+function _block_cache_type(::Type{Tv}, alg) where {Tv}
+    alg === nothing && return Nothing
+    return typeof(_block_cache(Tv, alg, 1))
+end
+
 function _dense_block_caches(
-        W::Vector{Matrix{Tv}}, sym::SymbolicAnalysis, dense_alg, dense_threshold::Int
-    ) where {Tv}
+        ::Type{Tv}, ::Type{C}, sym::SymbolicAnalysis, alg, dense_threshold::Int
+    ) where {Tv, C}
     nsuper = length(sym.sstart) - 1
-    bcaches = Vector{Any}(nothing, nsuper)
-    Tv <: LinearAlgebra.BlasFloat || return bcaches
-    assump = LinearSolve.OperatorAssumptions(true)
+    bcaches = Vector{Union{Nothing, C}}(nothing, nsuper)
+    alg === nothing && return bcaches
     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
-        )
+        bcaches[s] = _block_cache(Tv, alg, np)
     end
     return bcaches
 end
@@ -166,7 +187,29 @@ function _snlu_build(
     end
     V = M[sym.qf, sym.qf]
     Vt, Tpos = _transpose_map(V)
-    F = SupernodalLUFactor{Tv, Ti}(
+    dalg = _resolve_dense_alg(Tv, dense_alg)
+    # The dense block algorithm - and hence the cache type `C` - is a runtime
+    # choice (it follows LinearSolve's dense default, which depends on which
+    # packages are loaded).  Everything downstream of it is funnelled through
+    # this barrier so the factorization body is fully typed: the single
+    # dynamic call here is the only one in the construction path, and the
+    # resulting `SupernodalLUFactor{Tv, Ti, C}` is concrete.
+    return _snlu_assemble(
+        _block_cache_type(Tv, dalg), sym, A, M, ms, W, Z, rowsfac, V, Vt, Tpos,
+        dalg, dense_threshold, eps_pivot, threaded, check
+    )
+end
+
+function _snlu_assemble(
+        ::Type{C}, 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},
+        dalg, dense_threshold::Int, eps_pivot::Float64, threaded::Bool, check::Bool
+    ) where {C, Tv, Ti <: Integer}
+    n = sym.n
+    bcaches = _dense_block_caches(Tv, C, sym, dalg, dense_threshold)
+    F = SupernodalLUFactor{Tv, Ti, C}(
         sym, W, Z, collect(1:n), collect(1:n), rowsfac, 0, NaN, eps_pivot,
         A, Vector{Tv}(undef, n), Vector{Tv}(undef, 64),
         ms === nothing ? collect(1:n) : ms.rowperm,
@@ -178,7 +221,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
     )
     _build_vpos!(F, A)
     _factor_core!(F)

src/SupernodalLU/numeric.jl

@@ -44,7 +44,7 @@ Factorization object of [`snlu`](@ref).  Satisfies `A[F.p, F.q] ≈ F.L * F.U`
 with `L` unit lower triangular and `U` upper triangular.  `nperturbed(F)`
 reports how many pivots required static perturbation.
 """
-mutable struct SupernodalLUFactor{Tv, Ti <: Integer}
+mutable struct SupernodalLUFactor{Tv, Ti <: Integer, C}
     sym::SymbolicAnalysis
     W::Vector{Matrix{Tv}}            # supernode panels: [diag LU block; L21]
     Z::Vector{Matrix{Tv}}            # supernode panels: U12
@@ -77,9 +77,13 @@ 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); runtime-typed, accessed through the _cache_lu! barrier
-    bcaches::Vector{Any}
+    # Dense LinearSolve caches for the large diagonal blocks; `nothing` marks a
+    # block that uses the built-in kernel.  `C` is the concrete cache type,
+    # fixed by `(Tv, dense_alg)` alone (see `_resolve_dense_alg`), so the
+    # element type is a two-member union Julia union-splits into a static call
+    # in `_cache_lu!` and the whole factorization type is inferable from the
+    # inputs.  `C === Nothing` for element types that never cache a block.
+    bcaches::Vector{Union{Nothing, C}}
 end
 
 nperturbed(F::SupernodalLUFactor) = F.nperturbed
@@ -653,13 +657,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

@@ -12,7 +12,17 @@
 # factor-row space (update rows via `rowsfac`), the back solve in column space
 # (U12 columns are the un-permuted column ids in `rows`).
 
+# 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
@@ -28,7 +38,6 @@ function _solve_panels!(y::AbstractVector{Tv}, F::SupernodalLUFactor{Tv}) where
         xb = view(y, c1:c2)
         ldiv!(UnitLowerTriangular(view(Ws, 1:np, 1:np)), xb)
         if nu > 0
-            length(buf) < nu && resize!(buf, max(nu, 2 * length(buf)))
             t = view(buf, 1:nu)
             mul!(t, view(Ws, (np + 1):(np + nu), 1:np), xb)
             @simd for k in 1:nu
@@ -45,7 +54,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]]
@@ -58,6 +66,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
@@ -74,7 +83,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
@@ -91,7 +99,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]
@@ -114,6 +121,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]
@@ -140,6 +148,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

test/Core/supernodal_lu.jl

@@ -61,6 +61,42 @@ end
     @test (@allocated SNLU.snlu!(Fb, A)) == 0
 end
 
+@testset "block caches are concretely typed and inference-clean" begin
+    A = poisson2d(70)
+    F = SNLU.snlu(A)
+    @test count(!isnothing, F.bcaches) > 0
+    # the element type is a two-member union of Nothing and one concrete cache
+    # type - never Any - so the whole factorization type is inferable
+    @test eltype(F.bcaches) !== Any
+    @test isconcretetype(typeof(F).parameters[3])
+    @test isconcretetype(typeof(F))
+    b = randn(size(A, 1))
+    x = similar(b)
+    @inferred SNLU.solve!(x, F, b)
+    @test norm(A * x - b) <= 1.0e-11 * norm(b)
+    # element types that never cache a block collapse to Vector{Nothing}
+    T = spdiagm(0 => fill(big"4.0", 30), 1 => fill(big"-1.0", 29), -1 => fill(big"-1.0", 29))
+    Fb = SNLU.snlu(T)
+    @test typeof(Fb.bcaches) === Vector{Nothing}
+    @test isconcretetype(typeof(Fb))
+    # the LinearSolve cacheval prototype must have the same type as a real
+    # factorization, for every dense_alg setting
+    As = sparse(1.0I, 2000, 2000) + spdiagm(1 => fill(0.1, 1999))
+    bs = randn(2000)
+    for alg in (
+            SupernodalLUFactorization(),
+            SupernodalLUFactorization(dense_alg = LUFactorization()),
+            SupernodalLUFactorization(dense_alg = GenericLUFactorization()),
+        )
+        cache = init(LinearProblem(As, bs), alg)
+        sol = solve!(cache)
+        @test norm(As * sol.u - bs) <= 1.0e-10 * norm(bs)
+        cache.A = As
+        sol2 = solve!(cache)               # exercises cacheval reassignment
+        @test norm(As * sol2.u - bs) <= 1.0e-10 * norm(bs)
+    end
+end
+
 @testset "matching engages on zero/weak diagonals" begin
     n = 40
     P = sparse(collect(n:-1:1), collect(1:n), 2.0 .+ rand(n), n, n)

test/qa/Project.toml

@@ -11,7 +11,7 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 blis_jll = "6136c539-28a5-5bf0-87cc-b183200dce32"
 
 [sources]
-LinearSolve = {path = "../.."}
+LinearSolve = {path = "/home/crackauc/sandbox/tmp_20260717_160127_87884/LinearSolve.jl"}
 
 [compat]
 AllocCheck = "0.2"