SupernodalLU: refine only for perturbed pivots, drop residual-check allocation

ext/LinearSolveSparseArraysExt.jl

@@ -786,7 +786,8 @@ function SciMLBase.solve!(
     # so singularity surfaces as `Infeasible` instead of a silent `Success`.
     ok = all(isfinite, y)
     if ok && SNLU.nperturbed(F) > 0
-        r = F.A * y
+        r = F.ir_r                      # factor-owned residual buffer
+        LinearAlgebra.mul!(r, F.A, y)
         r .-= cache.b
         bn = LinearAlgebra.norm(cache.b)
         ok = LinearAlgebra.norm(r) <= 1.0e-6 * max(bn, floatmin(real(eltype(r))))

src/SupernodalLU/numeric.jl

@@ -77,8 +77,15 @@ 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
+    # 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.
     bcaches::Vector{Any}
 end
 

src/SupernodalLU/solve.jl

@@ -157,7 +157,11 @@ function _solve_once!(X::AbstractMatrix{Tv}, F::SupernodalLUFactor{Tv}, B::Abstr
     return X
 end
 
-_auto_refine(F::SupernodalLUFactor) = (F.nperturbed > 0 || F.matched) ? 3 : 0
+# Refinement exists to recover the accuracy lost to *perturbed* pivots.  MC64
+# matching improves conditioning rather than degrading it, so a matched factor
+# with no perturbed pivot needs none — refining there measured a 2.0-2.2x cost
+# per solve for a residual change of ~1.5e-15 -> 7e-16.
+_auto_refine(F::SupernodalLUFactor) = F.nperturbed > 0 ? 3 : 0
 
 """
     solve!(x, F::SupernodalLUFactor, b; refine=:auto) -> x

test/Core/supernodal_lu.jl

@@ -61,6 +61,44 @@ end
     @test (@allocated SNLU.snlu!(Fb, A)) == 0
 end
 
+@testset "auto refinement only for perturbed pivots" begin
+    # MC64 matching improves conditioning, so a matched factor with no
+    # perturbed pivot needs no refinement (refining there cost ~2.1x/solve).
+    n = 400
+    P = sparse(collect(n:-1:1), collect(1:n), 2.0 .+ rand(n), n, n)
+    Fm = SNLU.snlu(P)
+    @test Fm.matched
+    @test SNLU.nperturbed(Fm) == 0
+    @test SNLU._auto_refine(Fm) == 0
+    b = randn(n)
+    x = similar(b)
+    SNLU.solve!(x, Fm, b)
+    @test norm(P * x - b) <= 1.0e-12 * norm(b)
+    # a perturbed factor still refines
+    D = 10.0 .^ range(-9, 9; length = 200)
+    Ap = spdiagm(0 => D, 1 => fill(1.0e-9, 199), -1 => fill(1.0e-9, 199))
+    Fp = SNLU.snlu(Ap; matching = false, check = false)
+    @test SNLU.nperturbed(Fp) > 0
+    @test SNLU._auto_refine(Fp) == 3
+end
+
+@testset "residual check does not allocate a work vector" begin
+    # The post-solve residual check on a perturbed factor used to allocate a
+    # fresh n-vector (8n+104 = 2504 B at n=300); it now writes into the
+    # factor-owned buffer.  The bound (not == 0) tolerates the small
+    # `LinearSolution` each `solve!` returns, which the compiler elides only
+    # on some versions/platforms — it is still an order of magnitude below
+    # the old per-solve cost.
+    n = 300
+    Z = sprand(n, n, 0.02) + 1.0e-14I
+    cache = init(LinearProblem(Z, randn(n)), SupernodalLUFactorization())
+    solve!(cache)
+    resolve(c) = solve!(c)
+    resolve(cache)
+    resolve(cache)
+    @test (@allocated resolve(cache)) < 256
+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)