fix multiple solve handling

Author: ChrisRackauckas

ext/LinearSolveEnzymeExt.jl

@@ -1,6 +1,7 @@
 module LinearSolveEnzymeExt
 
 using LinearSolve
+using LinearSolve.LinearAlgebra
 isdefined(Base, :get_extension) ? (import Enzyme) : (import ..Enzyme)
 
 
@@ -60,9 +61,6 @@ end
 #   dA −= z y^T  
 #   dB += z, where  z = inv(A^T) dy
 function EnzymeCore.EnzymeRules.augmented_primal(config, func::Const{typeof(LinearSolve.solve!)}, ::Type{RT}, linsolve::EnzymeCore.Annotation{LP}; kwargs...) where {RT, LP <: LinearSolve.LinearCache}
-    @assert linsolve.val.isfresh
-    A_cache = copy(linsolve.val.A)
-
     res = func.val(linsolve.val; kwargs...)
 
     dres = if EnzymeRules.width(config) == 1
@@ -85,12 +83,12 @@ function EnzymeCore.EnzymeRules.augmented_primal(config, func::Const{typeof(Line
         (dr.u for dr in dres)
     end
 
-    cache = (A_cache, res, resvals)
+    cache = (res, resvals, linsolve.val)
     return EnzymeCore.EnzymeRules.AugmentedReturn(res, dres, cache)
 end
 
 function EnzymeCore.EnzymeRules.reverse(config, func::Const{typeof(LinearSolve.solve!)}, ::Type{RT}, cache, linsolve::EnzymeCore.Annotation{LP}; kwargs...) where {RT, LP <: LinearSolve.LinearCache}
-    A, y, dys = cache
+    y, dys, _linsolve = cache
 
     @assert !(typeof(linsolve) <: Const)
     @assert !(typeof(linsolve) <: Active)
@@ -112,11 +110,21 @@ function EnzymeCore.EnzymeRules.reverse(config, func::Const{typeof(LinearSolve.s
     end
 
     for (dA, db, dy) in zip(dAs, dbs, dys)
-        invprob = LinearSolve.LinearProblem(transpose(A), dy)
-        z = solve(invprob;
-            abstol = linsolve.val.abstol,
-            reltol = linsolve.val.reltol,
-            verbose = linsolve.val.verbose) 
+        z = if linsolve.cacheval isa Factorization
+            linsolve.cacheval' \ dy
+        elseif linsolve.cacheval isa Tuple && linsolve.cacheval[1] isa Factorization
+            linsolve.cacheval[1]' \ dy
+        elseif linsolve.alg isa AbstractKrylovSubspaceMethod
+            # Doesn't modify `A`, so it's safe to just reuse it
+            invprob = LinearSolve.LinearProblem(transpose(linsolve.A), dy)
+            solve(invprob;
+                abstol = linsolve.val.abstol,
+                reltol = linsolve.val.reltol,
+                verbose = linsolve.val.verbose,
+                isfresh = freshbefore)
+        else
+            error("Algorithm $(linsolve.alg) is currently not supported by Enzyme rules on LinearSolve.jl. Please open an issue on LinearSolve.jl detailing which algorithm is missing the adjoint handling")
+        end 
 
         dA .-= z * transpose(y)
         db .+= z

test/enzyme.jl

@@ -20,8 +20,8 @@ f(A, b1) # Uses BLAS
 
 Enzyme.autodiff(Reverse, f, Duplicated(copy(A), dA), Duplicated(copy(b1), db1))
 
-dA2 = FiniteDiff.finite_difference_gradient(x->f(x,b1, b2), copy(A))
-db12 = FiniteDiff.finite_difference_gradient(x->f(A,x, b2), copy(b1))
+dA2 = FiniteDiff.finite_difference_gradient(x->f(x,b1), copy(A))
+db12 = FiniteDiff.finite_difference_gradient(x->f(A,x), copy(b1))
 
 @test dA ≈ dA2
 @test db1 ≈ db12
@@ -33,6 +33,39 @@ b1 = rand(n);
 db1 = zeros(n);
 db12 = zeros(n);
 
-# This is not legal, all args need to be batch'd at the same size
-@test_broken Enzyme.autodiff(Reverse, f, Duplicated(copy(A), dA), BatchDuplicated(copy(b1), (db1, db12)))
-@test_broken Enzyme.autodiff(Reverse, f, BatchDuplicated(copy(A), (dA, dA2)), Duplicated(copy(b1), db1))
+@test_broken Enzyme.autodiff(Reverse, f, BatchDuplicated(copy(A), (dA, dA2)), BatchDuplicated(copy(b1), (db1, db12)))
+
+dA_2 = FiniteDiff.finite_difference_gradient(x->f(x,b1), copy(A))
+db1_2 = FiniteDiff.finite_difference_gradient(x->f(A,x), copy(b1))
+
+@test_broken dA ≈ dA_2
+@test_broken dA2 ≈ dA_2
+@test_broken db1 ≈ db1_2
+@test_broken db12 ≈ db1_2
+
+function f(A, b1, b2; alg = LUFactorization())
+    prob = LinearProblem(A, b1)
+
+    cache = init(prob, alg)
+    s1 = solve!(cache).u
+    cache.b = b2
+    s2 = solve!(cache).u
+    norm(s1 + s2)
+end
+
+A = rand(n, n);
+dA = zeros(n, n);
+b1 = rand(n);
+db1 = zeros(n);
+b2 = rand(n);
+db2 = zeros(n);
+
+Enzyme.autodiff(Reverse, f, Duplicated(copy(A), dA), Duplicated(copy(b1), db1), Duplicated(copy(b2), db2))
+
+dA2 = FiniteDiff.finite_difference_gradient(x->f(x,b1,b2), copy(A))
+db12 = FiniteDiff.finite_difference_gradient(x->f(A,x,b2), copy(b1))
+db22 = FiniteDiff.finite_difference_gradient(x->f(A,b1,x), copy(b2))
+
+@test dA ≈ dA2 atol=5e-5
+@test db1 ≈ db12
+@test db2 ≈ db22