Route reducible sparse-J WOperators to the sparse LHL solver (cost-model-gated default)

Project.toml

@@ -166,7 +166,7 @@ KernelAbstractions = "0.9.30"
 Krylov = "0.10"
 KrylovKit = "0.10"
 LAPACK_jll = "3"
-LHLFactorization = "2"
+LHLFactorization = "2.1.1"
 Libdl = "1.10"
 LinearAlgebra = "1.10"
 Markdown = "1.10"

src/LinearSolve.jl

@@ -8,7 +8,8 @@ import PrecompileTools
 using ArrayInterface: ArrayInterface
 # Explicit names, not the module: LinearSolve defines its own `LHLFactorization` (the
 # algorithm object) and `using LHLFactorization` would shadow it.
-using LHLFactorization: LHLWorkspace, lhl_reduce!, lhl_shift!, lhl_ldiv!, lhl_refine!
+using LHLFactorization: LHLWorkspace, lhl_reduce!, lhl_shift!, lhl_ldiv!, lhl_refine!,
+    lhl_ldivH!, lhl_refineH!, lhl, lhl!, lhl_isreduced, lhl_prefers_sparse
 using Base: Bool, convert, copyto!, adjoint, transpose, /, \, require_one_based_indexing
 using LinearAlgebra: LinearAlgebra, BlasInt, LU, Adjoint, BLAS, Bidiagonal, BunchKaufman,
     ColumnNorm, cond, Diagonal, Factorization, Hermitian, I, LAPACK, NoPivot,

src/adjoint_factorization.jl

@@ -90,11 +90,6 @@ for Alg in (
         ElementalJL,
         SpecializedLUFactorization,
         SpecializedQRFactorization,
-        # The LHL workspace is a Hessenberg reduction of J plus an LU of the shifted
-        # Hessenberg, not a `Factorization` of the system matrix; `Wᴴ = Z⁻ᴴ(I-γH)ᴴZᴴ` is
-        # solvable from the same pieces but is not implemented, so the adjoint reduces
-        # `Aᴴ` from scratch.
-        LHLFactorization,
     )
     @eval _adjoint_factorization_reuse(::Type{<:$Alg}) =
         _NoAdjointFactorizationReuse()
@@ -232,6 +227,59 @@ function _custom_adjoint_factorization_solve(
     return x
 end
 
+# `LHLFactorization` reuses its reduction on the reverse pass. The system is `σI + τJ` and
+# the cache already holds `J = ZHZ⁻¹` and the LU of the shifted Hessenberg for the current
+# shift, so `Aᴴ x = b` is `x = Z⁻ᴴ(σI+τH)⁻ᴴZᴴ b` — the same three `O(n²)` phases as the
+# forward solve, no refactorization (`lhl_ldivH!`), with the forward's iterative refinement
+# mirrored against `Aᴴ` (`lhl_refineH!`). Only the dense workspace has these adjoint
+# kernels; the block-triangular sparse factorization has none, so its adjoint is solved
+# against the assembled system with a fresh sparse LU rather than reused.
+_adjoint_factorization_reuse(::Type{<:LHLFactorization}) =
+    _CustomAdjointFactorizationReuse()
+
+_custom_can_reuse_adjoint_factorization(::LHLFactorization, c::LHLCache) =
+    c.ws isa LHLWorkspace
+_custom_can_reuse_adjoint_factorization(::LHLFactorization, ::Any) = false
+
+# `σI + τJ` assembled — the adjoint the reverse pass forms its residual against. A
+# `WOperator`'s own `adjoint` does not track `update_gamma!` (its forward `mul!` does, but
+# the adjoint keeps the γ it was built with), so refinement must not lean on `A'`; and the
+# generic reverse fallback cannot factorize an `adjoint(WOperator)` at all, so the sparse
+# branch assembles the matrix here too.
+_lhl_system_matrix(σ, τ, J) = σ * I + τ * J
+
+function _custom_adjoint_factorization_solve(alg::LHLFactorization, c::LHLCache, A, b)
+    (A isa AbstractMatrix || A isa WOperator) || return nothing
+    σ, τ = _lhl_shift_pair(A)
+    J = _lhl_jacobian(A)
+    if c.ws isa LHLWorkspace
+        ws = _lhl_sync!(c, A, alg, false)              # match the reduction and shift to A
+        ws.info == 0 || return nothing                 # singular shift: let the caller report it
+        T = promote_type(eltype(b), typeof(ws.σ))      # complex when the shift or b is
+        x = copyto!(similar(b, T), b)
+        M = alg.refine > 0 ? _lhl_system_matrix(σ, τ, J) : nothing
+        return _lhl_adjoint_reuse_solve!(x, M, b, ws, alg.refine)
+    end
+    # sparse LHL has no adjoint reduction: solve the adjoint of the assembled system.
+    return adjoint(_lhl_system_matrix(σ, τ, J)) \ b
+end
+
+function _lhl_adjoint_reuse_solve!(x::AbstractVector, M, b, ws, refine::Int)
+    lhl_ldivH!(x, ws)
+    refine > 0 && lhl_refineH!(x, M, b, ws, refine)
+    return x
+end
+
+# A batched (matrix) right-hand side solves column by column against the one reduction.
+function _lhl_adjoint_reuse_solve!(x::AbstractMatrix, M, b, ws, refine::Int)
+    for j in axes(x, 2)
+        xj = view(x, :, j)
+        lhl_ldivH!(xj, ws)
+        refine > 0 && lhl_refineH!(xj, M, view(b, :, j), ws, refine)
+    end
+    return x
+end
+
 """
     _adjoint_solve(cache::LinearCache, b)
 

src/default.jl

@@ -232,12 +232,16 @@ const LHL_DEFAULT_MIN_SIZE = 32
 # but left uninitialized the solve fails, and if it is initialized but never selected the
 # buffers are wasted. Both ask here.
 function _lhl_defaultable(A::WOperator, assump::OperatorAssumptions)
-    # Only a plain dense `J`. An operator `J` — a `MatrixOperator` in particular — is
-    # updated in place by `update_coefficients!`, which moves the numbers while leaving
-    # both the object identity and `jac_stale` untouched, so the reduction cannot tell it
-    # went stale and would silently answer with the previous Jacobian.
-    return assump.issq && A.J isa DenseMatrix && size(A, 1) >= LHL_DEFAULT_MIN_SIZE &&
-        _lhl_scalar_massmatrix(A.mass_matrix)
+    # Only a plain dense `J`, or a plain sparse `J` the sparse block-triangular solver is
+    # expected to win on (a reducible pattern — see `lhl_prefers_sparse`). An operator `J`
+    # — a `MatrixOperator` in particular — is updated in place by `update_coefficients!`,
+    # which moves the numbers while leaving both the object identity and `jac_stale`
+    # untouched, so the reduction cannot tell it went stale and would silently answer with
+    # the previous Jacobian, so those are excluded.
+    (assump.issq && size(A, 1) >= LHL_DEFAULT_MIN_SIZE && _lhl_scalar_massmatrix(A.mass_matrix)) ||
+        return false
+    A.J isa DenseMatrix && return true
+    return issparsematrixcsc(A.J) && lhl_prefers_sparse(A.J)
 end
 _lhl_defaultable(A, assump) = false
 

src/lhl.jl

@@ -115,10 +115,32 @@ function init_cacheval(
     (A isa AbstractMatrix || A isa WOperator) ||
         return LHLCache(LHLWorkspace{eltype(u)}(0), Nothing)
     J = _lhl_jacobian(A)
+    if _lhl_is_sparse(J)
+        # sparse `J`: the block-triangular sparse LHL of LHLFactorization's SparseArrays +
+        # PureKLU extension. `lhl(J)` analyzes and reduces; the returned factorization answers
+        # the same verbs as an `LHLWorkspace`.
+        F = lhl(J; shift = _lhl_shift_eltype(A, u), thread = _lhl_thread_bool(alg))
+        return LHLCache(F, typeof(J))
+    end
     ws = LHLWorkspace{eltype(J)}(size(A, 1); shift = _lhl_shift_eltype(A, u))
     return LHLCache(ws, typeof(J))
 end
 
+# `J` sparse enough to want the block-triangular sparse solver.
+_lhl_is_sparse(J) = issparsematrixcsc(J)
+_lhl_thread_bool(::LHLFactorization{T}) where {T} = T
+
+# The reduction step, dispatched on the workspace kind: a dense `LHLWorkspace` reduces in
+# place with the balance/thread the algorithm carries; the sparse factorization re-reduces
+# with `lhl!` (its analysis and per-block kernel choice are fixed at construction).
+_lhl_do_reduce!(ws::LHLWorkspace, J, alg::LHLFactorization) =
+    lhl_reduce!(ws, J, alg.balance, _lhl_thread(alg))
+_lhl_do_reduce!(F, J, ::LHLFactorization) = lhl!(F, J)
+
+# Size and reduced-state of the workspace, polymorphic over dense/sparse.
+_lhl_size1(ws::LHLWorkspace) = ws.n
+_lhl_size1(F) = size(F, 1)
+
 """
     _lhl_shift_eltype(A, u) -> Type
 
@@ -175,7 +197,7 @@ end
 # different `J` altogether, whose flag may already have been cleared by someone else.
 function _lhl_needs_reduce(c::LHLCache, A, isfresh::Bool)
     ws = c.ws
-    (ws.reduced && ws.n == size(A, 1)) || return true
+    (lhl_isreduced(ws) && _lhl_size1(ws) == size(A, 1)) || return true
     c.jac === _lhl_jacobian(A) || return true
     return _lhl_contents_moved(A, isfresh)
 end
@@ -191,7 +213,7 @@ _lhl_claim!(::AbstractMatrix) = nothing
 # The workspace's `setproperty!` forwards straight to `setfield!` without the conversion
 # Julia's default does, so a real `τ` cannot be stored into a complex shift (which is
 # exactly the real-J/complex-γ case) unless it is converted here.
-function _lhl_load_shift!(ws::LHLWorkspace, σ, τ)
+function _lhl_load_shift!(ws, σ, τ)
     lhl_shift!(ws, σ, τ)
     TG = typeof(ws.σ)
     ws.σ = convert(TG, σ)
@@ -205,7 +227,7 @@ function _lhl_sync!(c::LHLCache, A, alg::LHLFactorization, isfresh::Bool)
     fresh_reduction = _lhl_needs_reduce(c, A, isfresh)
     if fresh_reduction
         J = _lhl_jacobian(A)
-        lhl_reduce!(ws, J, alg.balance, _lhl_thread(alg))
+        _lhl_do_reduce!(ws, J, alg)
         c.jac = J
         _lhl_claim!(A)
     end

test/Core/lhl.jl

@@ -246,12 +246,86 @@ end
     @test bwd(ref1) < 1.0e-13
 end
 
-@testset "adjoint solve" begin
+@testset "adjoint solve reuses the reduction" begin
+    @test LinearSolve._adjoint_factorization_reuse(LHLFactorization) isa
+        LinearSolve._CustomAdjointFactorizationReuse
+
     n = 30
     J = randn(MersenneTwister(19), n, n)
     b = randn(MersenneTwister(20), n)
-    @test LinearSolve._adjoint_factorization_reuse(LHLFactorization) isa
-        LinearSolve._NoAdjointFactorizationReuse
+
+    # dense WOperator: the reverse pass solves Wᴴ x = b from the same reduction, and
+    # returning non-nothing proves it took the reuse path rather than the refactorizing
+    # fallback.
+    for γ in (0.4, 1.0e-3, 5.0)
+        cache = init(LinearProblem(wop(J, γ), b), LHLFactorization())
+        solve!(cache)
+        @test LinearSolve._custom_can_reuse_adjoint_factorization(cache.alg, cache.cacheval)
+        reused = LinearSolve._adjoint_factorization_solve(cache.alg, cache.cacheval, cache.A, b)
+        @test reused !== nothing
+        @test reused ≈ adjoint(dense(J, γ)) \ b rtol = 1.0e-9
+        @test LinearSolve._adjoint_solve(cache, b) ≈ adjoint(dense(J, γ)) \ b rtol = 1.0e-9
+    end
+
+    # a cheap re-shift keeps the adjoint current too
+    cache = init(LinearProblem(wop(J, 0.4), b), LHLFactorization())
+    solve!(cache)
+    update_gamma!(cache, 0.05)
+    @test LinearSolve._adjoint_solve(cache, b) ≈ adjoint(dense(J, 0.05)) \ b rtol = 1.0e-9
+
+    # complex γ on a real J: the adjoint keeps the real reduction and a complex shift
+    let γ = 0.2 + 0.3im, bc = randn(MersenneTwister(21), ComplexF64, n)
+        cache = init(
+            LinearProblem(WOperator{true}(I, γ, J, zeros(ComplexF64, n)), bc),
+            LHLFactorization()
+        )
+        solve!(cache)
+        @test LinearSolve._adjoint_solve(cache, bc) ≈ adjoint(J - I / γ) \ bc rtol = 1.0e-8
+    end
+
+    # a bare matrix reuses its reduction on the reverse pass as well
+    let A = randn(MersenneTwister(22), n, n)
+        cache = init(LinearProblem(A, b), LHLFactorization())
+        solve!(cache)
+        @test LinearSolve._adjoint_solve(cache, b) ≈ adjoint(A) \ b rtol = 1.0e-8
+    end
+
+    # a fully complex Jacobian
+    let Jc = randn(MersenneTwister(24), ComplexF64, n, n),
+            bc = randn(MersenneTwister(25), ComplexF64, n), γ = 0.3 + 0.1im
+
+        cache = init(
+            LinearProblem(wop(Jc, γ; u = zeros(ComplexF64, n)), bc),
+            LHLFactorization()
+        )
+        solve!(cache)
+        @test LinearSolve._adjoint_solve(cache, bc) ≈ adjoint(Jc - I / γ) \ bc rtol = 1.0e-8
+    end
+
+    # a batched (matrix) right-hand side solves column by column against the one reduction
+    let B = randn(MersenneTwister(26), n, 4), γ = 0.15
+        cache = init(LinearProblem(wop(J, γ), B[:, 1]), LHLFactorization())
+        solve!(cache)
+        X = LinearSolve._custom_adjoint_factorization_solve(cache.alg, cache.cacheval, cache.A, B)
+        @test X ≈ adjoint(dense(J, γ)) \ B rtol = 1.0e-8
+    end
+
+    # the block-triangular sparse LHL has no adjoint kernel, so it does not *reuse* the
+    # reduction (can_reuse is false), but the adjoint is still solved — against the
+    # assembled system with a fresh sparse LU
+    let Jsp = sparse([1, 2, 2, 3, 1], [1, 2, 3, 3, 3], [1.5, 2.0, 0.7, 1.1, 0.3], 3, 3),
+            bsp = randn(MersenneTwister(23), 3)
+
+        cache = init(LinearProblem(wop(Jsp, 0.1; u = zeros(3)), bsp), LHLFactorization())
+        solve!(cache)
+        @test !LinearSolve._custom_can_reuse_adjoint_factorization(cache.alg, cache.cacheval)
+        ref = adjoint(Matrix(Jsp) - I / 0.1) \ bsp
+        @test LinearSolve._adjoint_factorization_solve(cache.alg, cache.cacheval, cache.A, bsp) ≈
+            ref rtol = 1.0e-9
+        @test LinearSolve._adjoint_solve(cache, bsp) ≈ ref rtol = 1.0e-9
+    end
+
+    # still solves an assembled adjoint system when handed one directly
     @test solve(LinearProblem(adjoint(dense(J, 0.4)), b), LHLFactorization()).u ≈
         adjoint(dense(J, 0.4)) \ b rtol = 1.0e-9
 end
@@ -359,3 +433,64 @@ end
     cache.isfresh = true
     @test solve!(cache).u ≈ (Anew - I / γ) \ b rtol = 1.0e-6
 end
+
+# `J` sparse: LHLFactorization's SparseArrays + PureKLU extension (auto-loaded here, since
+# both are LinearSolve dependencies). `defaultalg` routes a *reducible* sparse-`J` WOperator
+# to LHL — where the block-triangular reduction can beat a sparse LU — and leaves one big
+# irreducible block to KLU.
+@testset "sparse Jacobian: solve, update_gamma!, cost-model-gated default" begin
+    assump = LinearSolve.OperatorAssumptions(true)
+
+    # a block upper triangular J with several irreducible blocks (reducible), permuted
+    function btf(sizes; rng = MersenneTwister(3))
+        m = sum(sizes)
+        I_ = Int[]; J_ = Int[]; V = Float64[]; off = 0
+        for bb in sizes
+            if bb == 1
+                push!(I_, off + 1); push!(J_, off + 1); push!(V, randn(rng))
+            else
+                for j in 1:bb, i in 1:bb
+                    (i == j || rand(rng) < 0.5) || continue
+                    push!(I_, off + i); push!(J_, off + j); push!(V, randn(rng))
+                end
+                for i in 1:(bb - 1)
+                    push!(I_, off + i + 1); push!(J_, off + i); push!(V, randn(rng))
+                end
+                push!(I_, off + 1); push!(J_, off + bb); push!(V, randn(rng))
+            end
+            off > 0 && (push!(I_, rand(rng, 1:off)); push!(J_, off + rand(rng, 1:bb)); push!(V, randn(rng)))
+            off += bb
+        end
+        Jm = sparse(I_, J_, V, m, m)
+        p = randperm(rng, m)
+        return Jm[p, p]
+    end
+
+    Js = btf([1, 3, 5, 8, 4, 20, 10])
+    ns = size(Js, 1)
+    bs = randn(MersenneTwister(5), ns)
+    Ws = wop(Js, 0.01; u = zeros(ns))
+    @test LinearSolve._lhl_defaultable(Ws, assump)        # reducible sparse -> default LHL
+    @test solve(LinearProblem(Ws, bs), LHLFactorization()).u ≈ dense(Js, 0.01) \ bs rtol = 1.0e-9
+
+    # cheap re-shift reuses the reduction
+    cache = init(LinearProblem(wop(Js, 0.01; u = zeros(ns)), bs), LHLFactorization())
+    solve!(cache)
+    for γ in (0.02, 1.0e-6, 5.0)
+        update_gamma!(cache, γ)
+        @test copy(solve!(cache).u) ≈ dense(Js, γ) \ bs rtol = 1.0e-9
+    end
+
+    # one big irreducible block: KLU's regime, not defaulted to LHL, but opt-in still solves
+    Jbig = btf([120])
+    @test !LinearSolve._lhl_defaultable(wop(Jbig, 0.01; u = zeros(120)), assump)
+    bb = randn(MersenneTwister(6), 120)
+    @test solve(LinearProblem(wop(Jbig, 0.01; u = zeros(120)), bb), LHLFactorization()).u ≈
+        dense(Jbig, 0.01) \ bb rtol = 1.0e-9
+
+    # complex γ on a real sparse J: real reduction, complex shift and solve
+    γc = 0.01 + 0.005im
+    Wc = WOperator{true}(I, γc, Js, zeros(ComplexF64, ns))
+    bc = randn(MersenneTwister(7), ComplexF64, ns)
+    @test solve(LinearProblem(Wc, bc), LHLFactorization()).u ≈ (Js - I / γc) \ bc rtol = 1.0e-8
+end