diff --git a/ext/LinearSolveSparseArraysExt.jl b/ext/LinearSolveSparseArraysExt.jl index 14457c407..12bf2969a 100644 --- a/ext/LinearSolveSparseArraysExt.jl +++ b/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 diff --git a/src/LinearSolve.jl b/src/LinearSolve.jl index 55b3d5683..044810f89 100644 --- a/src/LinearSolve.jl +++ b/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 diff --git a/src/SupernodalLU/interface.jl b/src/SupernodalLU/interface.jl index 376f00106..a15f0e33f 100644 --- a/src/SupernodalLU/interface.jl +++ b/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) diff --git a/src/SupernodalLU/numeric.jl b/src/SupernodalLU/numeric.jl index 72030ba69..693b8cd46 100644 --- a/src/SupernodalLU/numeric.jl +++ b/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 diff --git a/src/SupernodalLU/solve.jl b/src/SupernodalLU/solve.jl index 29da78d57..4c708a4e6 100644 --- a/src/SupernodalLU/solve.jl +++ b/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 diff --git a/src/SupernodalLU/symbolic.jl b/src/SupernodalLU/symbolic.jl index 1bad38d2a..10c9e1ff2 100644 --- a/src/SupernodalLU/symbolic.jl +++ b/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 diff --git a/test/Core/supernodal_lu.jl b/test/Core/supernodal_lu.jl index 29829444e..e6110718a 100644 --- a/test/Core/supernodal_lu.jl +++ b/test/Core/supernodal_lu.jl @@ -70,7 +70,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)) @@ -116,6 +116,45 @@ end @test (@allocated resolve(cache)) < 256 end +@testset "block caches are concretely typed and inference-clean" begin + A = poisson2d(70) + F = SNLU.snlu(A) + @test length(F.bcaches) > 0 + # the caches stay out of the factorization type, so it is fully concrete + # and inferable - which the sparse default solver depends on, since it + # holds a SupernodalLUFactor in one of its cacheval slots + @test typeof(F) === SNLU.SupernodalLUFactor{Float64, Int} + @test isconcretetype(typeof(F)) + @test count(!iszero, F.bcacheidx) == length(F.bcaches) + @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 isempty(Fb.bcaches) + @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 + 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) diff --git a/test/qa/allocations.jl b/test/qa/allocations.jl index f6ae16ec4..12b3da858 100644 --- a/test/qa/allocations.jl +++ b/test/qa/allocations.jl @@ -1,4 +1,4 @@ -using AllocCheck, LinearAlgebra, LinearSolve, Test +using AllocCheck, LinearAlgebra, LinearSolve, SparseArrays, Test if Sys.islinux() import LAPACK_jll, blis_jll @@ -62,3 +62,53 @@ if LinearSolve.appleaccelerate_isavailable() end end end + +# SupernodalLU: the triangular sweeps run entirely off buffers owned by the +# factorization, whose sizes are known from the symbolic analysis, so they +# carry no growth branch and AllocCheck can prove them allocation-free +# statically (rather than sampling `@allocated`). The user-facing `solve!` +# keeps a one-time scratch sizing, so it is asserted at runtime instead. +@check_allocs allocation_checked_supernodal_sweeps!(y, F) = + LinearSolve.SupernodalLU._solve_panels!(y, F) + +function poisson2d_qa(k) + n = k * k + Is = Int[]; Js = Int[]; V = Float64[] + idx(i, j) = (j - 1) * k + i + for j in 1:k, i in 1:k + c = idx(i, j) + push!(Is, c); push!(Js, c); push!(V, 4.0) + i > 1 && (push!(Is, c); push!(Js, idx(i - 1, j)); push!(V, -1.0)) + i < k && (push!(Is, c); push!(Js, idx(i + 1, j)); push!(V, -1.0)) + j > 1 && (push!(Is, c); push!(Js, idx(i, j - 1)); push!(V, -1.0)) + j < k && (push!(Is, c); push!(Js, idx(i, j + 1)); push!(V, -1.0)) + end + return sparse(Is, Js, V, n, n) +end + +@testset "SupernodalLU solve sweeps are provably allocation-free" begin + for A in ( + SparseArrays.spdiagm( + 0 => fill(4.0, 200), 1 => fill(-1.0, 199), -1 => fill(-1.0, 199) + ), + poisson2d_qa(30), # real panels: maxnu > 1 + ) + n = size(A, 1) + F = LinearSolve.SupernodalLU.snlu(A) + LinearSolve.SupernodalLU._ensure_panel_scratch!(F, 1) + y = ones(n) + allocation_checked_supernodal_sweeps!(y, F) # throws if it can allocate + @test true + # the full solve!, which owns the one-time sizing, is zero at runtime + b = ones(n) + x = similar(b) + LinearSolve.SupernodalLU.solve!(x, F, b; refine = 0) + @test @allocated(LinearSolve.SupernodalLU.solve!(x, F, b; refine = 0)) == 0 + @test norm(A * x - b) <= 1.0e-8 * norm(b) + # multi-RHS reuses the factor-owned scratch once sized + B = ones(n, 3) + X = similar(B) + LinearSolve.SupernodalLU.solve!(X, F, B; refine = 0) + @test @allocated(LinearSolve.SupernodalLU.solve!(X, F, B; refine = 0)) == 0 + end +end diff --git a/test/qa/jet.jl b/test/qa/jet.jl index fd995c9f6..f737219e5 100644 --- a/test/qa/jet.jl +++ b/test/qa/jet.jl @@ -256,3 +256,44 @@ end end end end + +@testset "JET Tests for SupernodalLU" begin + # The dense block caches are concretely typed, so the repeated-use paths - + # solve and numeric refactorization, i.e. the ODE/Newton workload - must be + # free of runtime dispatch. Analysis is scoped to the solver module, as + # elsewhere in this file, so Base error-path noise (show/AssertionError + # reachable from sparse indexing) does not mask our own code. + SNLU_MOD = LinearSolve.SupernodalLU + n = 60 + A_snlu = SparseArrays.spdiagm( + 0 => fill(4.0, n), 1 => fill(-1.0, n - 1), -1 => fill(-1.0, n - 1) + ) + b_snlu = ones(n) + F_snlu = SNLU_MOD.snlu(A_snlu) + x_snlu = similar(b_snlu) + + # The solve path never touches the block caches, so it must be free of + # runtime dispatch. + JET.@test_opt target_modules = (SNLU_MOD,) SNLU_MOD.solve!(x_snlu, F_snlu, b_snlu) + JET.@test_opt target_modules = (SNLU_MOD,) SNLU_MOD._solve_panels!(x_snlu, F_snlu) + + # Factorization carries exactly one dispatch per *cached* supernode: the + # `_cache_lu!` barrier, whose cache argument is deliberately untyped (see + # the `bcaches` field docs - putting the cache type in the factorization + # type breaks `@inferred init(prob, nothing)` for sparse matrices, because + # the sparse default holds a SupernodalLUFactor in its cacheval). Tens of + # calls against hundreds of ms of GEMM; the alternative costs far more. + JET.@test_opt target_modules = (SNLU_MOD,) SNLU_MOD.snlu!(F_snlu, A_snlu) broken = true + JET.@test_opt target_modules = (SNLU_MOD,) SNLU_MOD.snlu(A_snlu) broken = true + + # The factorization object itself is concrete, and its block-cache vector + # is a two-member union rather than Any. + @test isconcretetype(typeof(F_snlu)) + @test typeof(F_snlu) === SNLU_MOD.SupernodalLUFactor{Float64, Int} + + # The factorization type itself is fully concrete and inferable from its + # inputs, which is what the sparse default solver depends on. + @test isconcretetype( + Core.Compiler.return_type(SNLU_MOD.snlu, Tuple{typeof(A_snlu)}) + ) +end