diff --git a/ext/LinearSolveForwardDiffExt.jl b/ext/LinearSolveForwardDiffExt.jl index 97d03d7bd..561d06e7a 100644 --- a/ext/LinearSolveForwardDiffExt.jl +++ b/ext/LinearSolveForwardDiffExt.jl @@ -56,8 +56,9 @@ LinearSolve.@concrete mutable struct DualLinearCache{DT} primal_u_cache primal_b_cache - # Cache validity flag for RHS precalculation optimization - rhs_cache_valid + # Cache validity flags for when partials of A or b changes + A_partials_valid + b_partials_valid dual_A dual_b @@ -94,15 +95,13 @@ function linearsolve_forwarddiff_solve!(cache::DualLinearCache, alg, args...; kw rhs_list = cache.rhs_list - # Update cached partials lists if cache is invalid - if !cache.rhs_cache_valid - if !isnothing(∂_A) - update_partials_list!(∂_A, cache.partials_A_list) - end - if !isnothing(∂_b) - update_partials_list!(∂_b, cache.partials_b_list) - end - cache.rhs_cache_valid = true + if !cache.A_partials_valid && !isnothing(∂_A) + update_partials_list!(∂_A, cache.partials_A_list) + cache.A_partials_valid = true + end + if !cache.b_partials_valid && !isnothing(∂_b) + update_partials_list!(∂_b, cache.partials_b_list) + cache.b_partials_valid = true end A_list = cache.partials_A_list @@ -177,11 +176,13 @@ function xp_linsolve_rhs!( ∂_b::Union{<:Partials, <:AbstractArray{<:Partials}}, cache::DualLinearCache ) - # Update cached partials lists if cache is invalid - if !cache.rhs_cache_valid + if !cache.A_partials_valid update_partials_list!(∂_A, cache.partials_A_list) + cache.A_partials_valid = true + end + if !cache.b_partials_valid update_partials_list!(∂_b, cache.partials_b_list) - cache.rhs_cache_valid = true + cache.b_partials_valid = true end A_list = cache.partials_A_list @@ -201,10 +202,9 @@ function xp_linsolve_rhs!( ∂_b::Nothing, cache::DualLinearCache ) - # Update cached partials list for A if cache is invalid - if !cache.rhs_cache_valid + if !cache.A_partials_valid update_partials_list!(∂_A, cache.partials_A_list) - cache.rhs_cache_valid = true + cache.A_partials_valid = true end A_list = cache.partials_A_list @@ -222,10 +222,9 @@ function xp_linsolve_rhs!( cache::DualLinearCache ) - # Update cached partials list for b if cache is invalid - if !cache.rhs_cache_valid + if !cache.b_partials_valid update_partials_list!(∂_b, cache.partials_b_list) - cache.rhs_cache_valid = true + cache.b_partials_valid = true end b_list = cache.partials_b_list @@ -380,6 +379,7 @@ function __dual_init( similar(non_partial_cache.u), # primal_u_cache similar(new_b), # primal_b_cache true, # Cache is initially valid + true, A, b, dual_u_init @@ -495,8 +495,7 @@ function setA!(dc::DualLinearCache, A) partial_vals!(getfield(dc, :partials_A), A) end - # Invalidate cache (if setting A or b) - return setfield!(dc, :rhs_cache_valid, false) + return setfield!(dc, :A_partials_valid, false) end function setb!(dc::DualLinearCache, b) # Put the Dual-stripped versions in the LinearCache @@ -510,8 +509,7 @@ function setb!(dc::DualLinearCache, b) partial_vals!(getfield(dc, :partials_b), b) end - # Invalidate cache (if setting A or b) - return setfield!(dc, :rhs_cache_valid, false) + return setfield!(dc, :b_partials_valid, false) end function setu!(dc::DualLinearCache{DT}, u) where {DT} # Put the Dual-stripped versions in the LinearCache diff --git a/test/Core/forwarddiff_overloads.jl b/test/Core/forwarddiff_overloads.jl index 373f60e20..182023b0a 100644 --- a/test/Core/forwarddiff_overloads.jl +++ b/test/Core/forwarddiff_overloads.jl @@ -571,3 +571,38 @@ end extract_partials(v) = [collect(ForwardDiff.partials(x)) for x in v] @test all(((a, b),) -> a ≈ b, zip(extract_partials(sol.u), extract_partials(ref))) end + +# The DualLinearCache tracks partials-list validity for A and b independently, +# so mutating only one side does not force the other's partials to be recomputed +# (relevant e.g. in an ODE where A is fixed while b changes, and vice versa). +@testset "DualLinearCache separate A/b partials validity" begin + A, b = h([ForwardDiff.Dual(10.0, 1.0, 0.0), ForwardDiff.Dual(10.0, 0.0, 1.0)]) + cache = init(LinearProblem(A, b), LUFactorization()) + + # Both lists start valid (populated lazily on first solve). + @test getfield(cache, :A_partials_valid) + @test getfield(cache, :b_partials_valid) + + # Mutating only b invalidates b's list and leaves A's untouched. + _, new_b = h([ForwardDiff.Dual(5.0, 1.0, 0.0), ForwardDiff.Dual(5.0, 0.0, 1.0)]) + cache.b = new_b + @test getfield(cache, :A_partials_valid) + @test !getfield(cache, :b_partials_valid) + + # Solving revalidates both, and the result still matches the reference. + x_p = solve!(cache) + @test getfield(cache, :A_partials_valid) + @test getfield(cache, :b_partials_valid) + @test ≈(x_p, A \ new_b, rtol = 1.0e-9) + + # Symmetrically, mutating only A invalidates A's list and leaves b's untouched. + new_A, _ = h([ForwardDiff.Dual(2.0, 1.0, 0.0), ForwardDiff.Dual(2.0, 0.0, 1.0)]) + cache.A = new_A + @test !getfield(cache, :A_partials_valid) + @test getfield(cache, :b_partials_valid) + + x_p = solve!(cache) + @test getfield(cache, :A_partials_valid) + @test getfield(cache, :b_partials_valid) + @test ≈(x_p, new_A \ new_b, rtol = 1.0e-9) +end