From 0f1ce09571153b190669f1f1083600a700dc375f Mon Sep 17 00:00:00 2001 From: Herman Sletmoen Date: Sun, 28 Jun 2026 14:09:17 +0200 Subject: [PATCH 1/4] Cache inner LinearCache in DualLinearCache in direct dual path --- ext/LinearSolveForwardDiffExt.jl | 52 ++++++++++++++++++++++++-------- 1 file changed, 40 insertions(+), 12 deletions(-) diff --git a/ext/LinearSolveForwardDiffExt.jl b/ext/LinearSolveForwardDiffExt.jl index 4acc7e9ad..e201684ee 100644 --- a/ext/LinearSolveForwardDiffExt.jl +++ b/ext/LinearSolveForwardDiffExt.jl @@ -63,6 +63,9 @@ LinearSolve.@concrete mutable struct DualLinearCache{DT} dual_A dual_b dual_u + + # Cached LinearCache for direct Dual path (nothing for split path algorithms) + dual_linear_cache end function linearsolve_forwarddiff_solve!(cache::DualLinearCache, alg, args...; kwargs...) @@ -367,6 +370,20 @@ function __dual_init( ArrayInterface.restructure(non_partial_cache.u, zeros(dual_type, length(non_partial_cache.u))) end + # For algorithms taking the direct Dual path, use __init to create a regular LinearCache + # (bypasses ForwardDiff extension) then solve! on that cache directly with the dual values + # Promote b to Dual so that dc.b is typed correctly for later dc.b = dual_b + # assignments in _solve_direct_dual! (b may be plain Float64 when only A is Dual). + if _use_direct_dual_solve(alg) + dual_b_init = eltype(b) <: ForwardDiff.Dual ? b : dual_type.(b) + dual_linear_cache_init = __init( + LinearProblem(A, dual_b_init), alg; + alias, abstol, reltol, maxiters, verbose, Pl, Pr, assumptions, sensealg, kwargs... + ) + else + dual_linear_cache_init = nothing + end + return DualLinearCache{dual_type}( non_partial_cache, ∂_A, @@ -382,7 +399,8 @@ function __dual_init( true, A, b, - dual_u_init + dual_u_init, + dual_linear_cache_init ) end @@ -444,6 +462,8 @@ end # Direct solve path for algorithms that can work with Dual numbers directly # This avoids the primal/partials separation overhead +# The inner dual LinearCache is created eagerly in __dual_init and reused here, +# mirroring how the split path reuses the primal factorisation across RHS vectors. function _solve_direct_dual!( cache::DualLinearCache{DT}, alg, args...; kwargs... ) where {DT <: ForwardDiff.Dual} @@ -451,19 +471,28 @@ function _solve_direct_dual!( dual_A = getfield(cache, :dual_A) dual_b = getfield(cache, :dual_b) - # When the duals are only in A, b is stored as a plain array. Promote it to - # the dual type: `__init` takes the solution eltype from `b`, and the - # factorization/solve kernels need matching element types. + # When only A carries duals, b is a plain array — promote it so the + # factorisation kernel sees a uniform element type. if eltype(dual_b) != DT dual_b = DT.(dual_b) end - # Use __init to create a regular LinearCache (bypasses ForwardDiff extension) - # then solve! on that cache directly with the dual values - dual_prob = LinearProblem(dual_A, dual_b) - dual_cache = __init(dual_prob, getfield(cache, :linear_cache).alg, args...; kwargs...) + # Get regular LinearCache prepared in __dual__init + linear_cache = getfield(cache, :linear_cache) + dual_cache = getfield(cache, :dual_linear_cache) + + # Update A (and trigger re-factorisation) when the outer primal cache signals + # that A has changed via its isfresh flag, which is set by setA!. + if linear_cache.isfresh + dual_cache.A = dual_A + end + dual_cache.b = dual_b + + # solve! on the regular LinearCache directly with the dual values (bypasses ForwardDiff extension) dual_sol = SciMLBase.solve!(dual_cache) + setfield!(linear_cache, :isfresh, false) + # Update the cache if getfield(cache, :dual_u) isa AbstractArray getfield(cache, :dual_u) .= dual_sol.u @@ -472,13 +501,12 @@ function _solve_direct_dual!( end # Also update the primal cache for consistency - primal_u = nodual_value.(dual_sol.u) - if getfield(cache, :linear_cache).u isa AbstractArray - getfield(cache, :linear_cache).u .= primal_u + if linear_cache.u isa AbstractArray + linear_cache.u .= nodual_value.(dual_sol.u) end return SciMLBase.build_linear_solution( - getfield(cache, :linear_cache).alg, getfield(cache, :dual_u), dual_sol.resid, cache; + linear_cache.alg, getfield(cache, :dual_u), dual_sol.resid, cache; dual_sol.retcode, dual_sol.iters, dual_sol.stats ) end From 53a2e0eee49f809bb1b3cf894e20b34356f20c62 Mon Sep 17 00:00:00 2001 From: Herman Sletmoen Date: Sun, 28 Jun 2026 18:55:11 +0200 Subject: [PATCH 2/4] Remove runtime _use_direct_dual_solve check that prevents compile-time inference of whether dual_linear_cache field is nothing --- ext/LinearSolveForwardDiffExt.jl | 4 ---- 1 file changed, 4 deletions(-) diff --git a/ext/LinearSolveForwardDiffExt.jl b/ext/LinearSolveForwardDiffExt.jl index e201684ee..1825ea73c 100644 --- a/ext/LinearSolveForwardDiffExt.jl +++ b/ext/LinearSolveForwardDiffExt.jl @@ -422,10 +422,6 @@ function _use_direct_dual_solve(alg) alg isa LinearSolve.PureKLUFactorization end -function _use_direct_dual_solve(alg::DefaultLinearSolver) - return alg.alg === DefaultAlgorithmChoice.GenericLUFactorization -end - function SciMLBase.solve!( cache::DualLinearCache{DT}, alg::SciMLLinearSolveAlgorithm, args...; kwargs... ) where { From 2b44c20db1b49d49c714ac29f10636e721cd1ea2 Mon Sep 17 00:00:00 2001 From: Herman Sletmoen Date: Sun, 28 Jun 2026 22:06:30 +0200 Subject: [PATCH 3/4] Copy A/b if necessary in direct dual path --- ext/LinearSolveForwardDiffExt.jl | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/ext/LinearSolveForwardDiffExt.jl b/ext/LinearSolveForwardDiffExt.jl index 1825ea73c..73441dbfd 100644 --- a/ext/LinearSolveForwardDiffExt.jl +++ b/ext/LinearSolveForwardDiffExt.jl @@ -2,7 +2,7 @@ module LinearSolveForwardDiffExt using LinearSolve using LinearSolve: SciMLLinearSolveAlgorithm, __init, LinearVerbosity, DefaultLinearSolver, - DefaultAlgorithmChoice, defaultalg, reinit! + DefaultAlgorithmChoice, defaultalg, reinit!, default_alias_A using LinearAlgebra using SparseArrays using ForwardDiff @@ -480,9 +480,9 @@ function _solve_direct_dual!( # Update A (and trigger re-factorisation) when the outer primal cache signals # that A has changed via its isfresh flag, which is set by setA!. if linear_cache.isfresh - dual_cache.A = dual_A + dual_cache.A = default_alias_A(alg, dual_A, dual_b) ? dual_A : copy(dual_A) end - dual_cache.b = dual_b + dual_cache.b .= dual_b # solve! on the regular LinearCache directly with the dual values (bypasses ForwardDiff extension) dual_sol = SciMLBase.solve!(dual_cache) From 55e7f19bb7dbdf87beb788935123b8c621d2650a Mon Sep 17 00:00:00 2001 From: Herman Sletmoen Date: Sun, 28 Jun 2026 19:13:13 +0200 Subject: [PATCH 4/4] Test that PureKLU direct dual path reuses inner LinearCache --- test/Core/forwarddiff_overloads.jl | 35 ++++++++++++++++++++++++++++++ 1 file changed, 35 insertions(+) diff --git a/test/Core/forwarddiff_overloads.jl b/test/Core/forwarddiff_overloads.jl index 978a4faa3..9b6850da7 100644 --- a/test/Core/forwarddiff_overloads.jl +++ b/test/Core/forwarddiff_overloads.jl @@ -78,6 +78,7 @@ prob = LinearProblem(plain_A, b) # Sanity: the genuinely-cheap-in-dual algorithms are still on the direct path. @test ext._use_direct_dual_solve(GenericLUFactorization()) @test ext._use_direct_dual_solve(SpecializedLUFactorization()) + @test ext._use_direct_dual_solve(PureKLUFactorization()) end # Overload Dense @@ -646,3 +647,37 @@ end cache.A = A @test ≈(solve!(cache), Matrix(A) \ b, rtol = 1.0e-9) end + +@testset "PureKLU direct dual path reuses inner LinearCache" begin + # The direct dual path (PureKLUFactorization can handle Dual arithmetic + # natively) pre-creates an inner LinearCache in DualLinearCache.dual_linear_cache + # at init time and reuses it on every solve!, only re-factorising when + # linear_cache.isfresh is set. Verify that the inner cache object identity + # is preserved across successive solves with different A and b. + A1 = sparse([1, 2], [1, 2], [ForwardDiff.Dual(1.0, 10.0, 0.0), ForwardDiff.Dual(2.0, 0.0, 20.0)], 2, 2) + b1 = [ForwardDiff.Dual(3.0, 1.0, 0.0), ForwardDiff.Dual(4.0, 0.0, 1.0)] + prob = LinearProblem(A1, b1) + cache = init(prob, PureKLUFactorization()) + + # The inner dual cache must be pre-created at init time (not nothing). + inner = getfield(cache, :dual_linear_cache) + @test inner !== nothing + + sol1 = solve!(cache) + @test ≈(sol1, Matrix(A1) \ b1, rtol = 1.0e-9) + + # After mutating A the inner cache object must be the same — only its + # contents change, no new allocation. + A2 = sparse([1, 2], [1, 2], [ForwardDiff.Dual(3.0, 30.0, 0.0), ForwardDiff.Dual(4.0, 0.0, 40.0)], 2, 2) + cache.A = A2 + sol2 = solve!(cache) + @test getfield(cache, :dual_linear_cache) === inner + @test ≈(sol2, Matrix(A2) \ b1, rtol = 1.0e-9) + + # Same after mutating only b. + b2 = [ForwardDiff.Dual(5.0, 2.0, 0.0), ForwardDiff.Dual(6.0, 0.0, 2.0)] + cache.b = b2 + sol3 = solve!(cache) + @test getfield(cache, :dual_linear_cache) === inner + @test ≈(sol3, Matrix(A2) \ b2, rtol = 1.0e-9) +end