Store the dense block caches sparsely, dropping the Union element type

Only the blocks at or above dense_threshold get a LinearSolve cache;
everything narrower uses the built-in kernel. That 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 - so a
full-length Vector{Union{Nothing,C}} was 99 % empty and forced a
nullability check on every lookup.

bcaches now holds only the blocks that have a cache, with
bcacheidx[s] indexing into it (0 = use the built-in kernel). The element
type is plainly concrete rather than a two-member union, the lookup is an
integer test, and the factorization type stays inferable from its inputs
(return_type(snlu, ...) concrete). For element types that never cache a
block it degenerates to an empty Vector{Nothing} with an all-zero index.

Caching every supernode instead was considered and rejected on the
numbers: a default-solver 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 that motivates
PANEL_BLAS_CUTOFF in the solve sweeps.

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

src/SupernodalLU/interface.jl

@@ -139,14 +139,16 @@ function _dense_block_caches(
         ::Type{Tv}, ::Type{C}, sym::SymbolicAnalysis, alg, dense_threshold::Int
     ) where {Tv, C}
     nsuper = length(sym.sstart) - 1
-    bcaches = Vector{Union{Nothing, C}}(nothing, nsuper)
-    C === Nothing && return bcaches
+    bcaches = C[]
+    bcacheidx = zeros(Int, nsuper)
+    C === Nothing && return bcaches, bcacheidx
     for s in 1:nsuper
         np = sym.sstart[s + 1] - sym.sstart[s]
         np >= dense_threshold || continue
-        bcaches[s] = _block_cache(Tv, alg, np)
+        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)
@@ -188,7 +190,7 @@ function _snlu_assemble(
         check::Bool
     ) where {C, Tv, Ti <: Integer}
     n = sym.n
-    bcaches = _dense_block_caches(Tv, C, sym, dense_alg, dense_threshold)
+    bcaches, bcacheidx = _dense_block_caches(Tv, C, sym, dense_alg, 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),
@@ -201,7 +203,7 @@ function _snlu_assemble(
         Matrix{Tv}(undef, n, 0),
         Vector{Tv}(undef, n), Vector{Tv}(undef, n), Vector{Tv}(undef, n),
         threaded,
-        bcaches
+        bcaches, bcacheidx
     )
     _build_vpos!(F, A)
     _factor_core!(F)

src/SupernodalLU/numeric.jl

@@ -77,13 +77,18 @@ mutable struct SupernodalLUFactor{Tv, Ti <: Integer, C}
     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 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}}
+    # 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".  That
+    # keeps the element type concrete (`C`, not `Union{Nothing,C}`) and avoids
+    # a mostly-empty full-length vector.  `C` is fixed by `(Tv, dense_alg)`
+    # alone, so the whole factorization type is inferable from the inputs;
+    # `C === Nothing` for element types that never cache a block.
+    bcaches::Vector{C}
+    bcacheidx::Vector{Int}
 end
 
 nperturbed(F::SupernodalLUFactor) = F.nperturbed
@@ -234,9 +239,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

test/Core/supernodal_lu.jl

@@ -53,7 +53,7 @@ end
     # `LinearSolution` object each dense block cache's `solve!` returns —
     # a fixed ~64 B per cached supernode, independent of problem size.
     @test (@allocated SNLU.solve!(x, F, b)) == 0
-    ncache = count(!isnothing, F.bcaches)
+    ncache = length(F.bcaches)
     @test (@allocated SNLU.snlu!(F, A)) <= 128 * max(ncache, 1)
     # with the dense caches disabled the built-in kernel is fully allocation-free
     Fb = SNLU.snlu(A; dense_threshold = typemax(Int))
@@ -64,10 +64,12 @@ end
 @testset "block caches are concretely typed and inference-clean" begin
     A = poisson2d(70)
     F = SNLU.snlu(A)
-    @test count(!isnothing, F.bcaches) > 0
+    @test length(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(eltype(F.bcaches))
+    @test count(!iszero, F.bcacheidx) == length(F.bcaches)
     @test isconcretetype(typeof(F).parameters[3])
     @test isconcretetype(typeof(F))
     b = randn(size(A, 1))
@@ -78,6 +80,7 @@ end
     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 all(iszero, Fb.bcacheidx)
     @test isconcretetype(typeof(Fb))
     # the LinearSolve cacheval prototype must have the same type as a real
     # factorization, for every dense_alg setting

test/qa/jet.jl

@@ -280,6 +280,7 @@ end
     # is a two-member union rather than Any.
     @test isconcretetype(typeof(F_snlu))
     @test eltype(F_snlu.bcaches) !== Any
+    @test isconcretetype(eltype(F_snlu.bcaches))
     @test isconcretetype(typeof(F_snlu).parameters[3])
 
     # Construction is inferable too: the block caches use LinearSolve's own