Cache dense blocks with LinearSolve's own default solver

Follow-up on the concrete-typing commit: instead of resolving a dense
algorithm ourselves and parameterising on the resulting cache type, each
cached block now uses LinearSolve's default solver directly
(`init(prob, nothing)`).

Its cache type is
`LinearCache{..., DefaultLinearSolver, DefaultLinearSolverInit{...}}`,
which holds a slot for every candidate algorithm and picks between them
with a runtime enum - so the type depends only on the element type, not
on which algorithm the default happens to select on this machine. That
makes the whole factorization type inferable from its inputs:
`Core.Compiler.return_type(snlu, Tuple{SparseMatrixCSC{Float64,Int}})`
is now concrete, and JET reports 0 runtime dispatches for construction
as well as for solve! and snlu! (construction was 3 before).

It is also simply what "use the LinearSolve default" should mean: the
bespoke `_resolve_dense_alg` probe is gone.

The extra slots cost nothing measurable - 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.

The factorization behind a block cache is pulled out of the slot the
enum names, written as a literal-symbol chain so every branch stays
concrete. src/SupernodalLU is now included after src/default.jl, since
it references DefaultLinearSolver{,Init}.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

Author: ChrisRackauckas

src/LinearSolve.jl

@@ -409,7 +409,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")
@@ -426,6 +425,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,29 +110,16 @@ 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.
+# 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)),
@@ -142,10 +129,9 @@ function _block_cache(::Type{Tv}, alg, np::Int) where {Tv}
     )
 end
 
-# Concrete element type of `bcaches` for this element type and algorithm.
-# `Nothing` when no block will ever be cached (non-BLAS element types).
+# Concrete element type of `bcaches`; `Nothing` when no block is ever cached.
 function _block_cache_type(::Type{Tv}, alg) where {Tv}
-    alg === nothing && return Nothing
+    Tv <: LinearAlgebra.BlasFloat || return Nothing
     return typeof(_block_cache(Tv, alg, 1))
 end
 
@@ -154,7 +140,7 @@ function _dense_block_caches(
     ) where {Tv, C}
     nsuper = length(sym.sstart) - 1
     bcaches = Vector{Union{Nothing, C}}(nothing, nsuper)
-    alg === nothing && return bcaches
+    C === Nothing && return bcaches
     for s in 1:nsuper
         np = sym.sstart[s + 1] - sym.sstart[s]
         np >= dense_threshold || continue
@@ -187,16 +173,9 @@ function _snlu_build(
     end
     V = M[sym.qf, sym.qf]
     Vt, Tpos = _transpose_map(V)
-    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
+        _block_cache_type(Tv, dense_alg), sym, A, M, ms, W, Z, rowsfac, V, Vt,
+        Tpos, dense_alg, dense_threshold, eps_pivot, threaded, check
     )
 end
 
@@ -205,10 +184,11 @@ function _snlu_assemble(
         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
+        dense_alg, 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)
+    bcaches = _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),

src/SupernodalLU/numeric.jl

@@ -256,7 +256,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
@@ -274,6 +274,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}

test/qa/jet.jl

@@ -282,10 +282,12 @@ end
     @test eltype(F_snlu.bcaches) !== Any
     @test isconcretetype(typeof(F_snlu).parameters[3])
 
-    # Construction resolves the dense block algorithm at runtime (it follows
-    # LinearSolve's dense default, which depends on what is loaded), so the
-    # cache type - and hence the factorization type - is only known then.  The
-    # choice is funnelled through one function barrier, so this is a handful of
-    # dispatches in a one-time path, not in the repeated ones asserted above.
-    JET.@test_opt target_modules = (SNLU_MOD,) SNLU_MOD.snlu(A_snlu) broken = true
+    # Construction is inferable too: the block caches use LinearSolve's own
+    # default solver, whose cache type holds a slot per candidate algorithm and
+    # selects with a runtime enum, so the type does not depend on which
+    # algorithm the default picks on this machine.
+    JET.@test_opt target_modules = (SNLU_MOD,) SNLU_MOD.snlu(A_snlu)
+    @test isconcretetype(
+        Core.Compiler.return_type(SNLU_MOD.snlu, Tuple{typeof(A_snlu)})
+    )
 end