Use separate cache flags for A and b in ForwardDiff extension (rebased #1054 + test) (#1062)

* Use separate cache flags for A and b (faster when only one changes)

* Test separate A/b partials validity flags in DualLinearCache

Verify that mutating only A or only b invalidates that side's partials
list independently, that solving revalidates both, and that the result
stays correct across one-sided updates.

Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>

---------

Co-authored-by: Herman Sletmoen <hermansletmoen@gmail.com>
Co-authored-by: ChrisRackauckas-Claude <accounts@chrisrackauckas.com>

Author: 3 people

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

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