SupernodalLU: sparse block-cache storage, AllocCheck/JET coverage, threaded race fix

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