Add forward enzyme rules for init and solve

Author: sharanry

Project.toml

@@ -66,6 +66,7 @@ DocStringExtensions = "0.9"
 EnumX = "1"
 EnzymeCore = "0.6"
 FastLapackInterface = "2"
+EnzymeTestUtils = "0.1"
 GPUArraysCore = "0.1"
 HYPRE = "1.4.0"
 InteractiveUtils = "1.6"
@@ -96,6 +97,7 @@ julia = "1.9"
 BandedMatrices = "aae01518-5342-5314-be14-df237901396f"
 BlockDiagonals = "0a1fb500-61f7-11e9-3c65-f5ef3456f9f0"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+EnzymeTestUtils = "12d8515a-0907-448a-8884-5fe00fdf1c5a"
 FiniteDiff = "6a86dc24-6348-571c-b903-95158fe2bd41"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 HYPRE = "b5ffcf37-a2bd-41ab-a3da-4bd9bc8ad771"
@@ -114,4 +116,4 @@ SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "IterativeSolvers", "InteractiveUtils", "JET", "KrylovKit", "Pkg", "Random", "SafeTestsets", "MultiFloats", "ForwardDiff", "HYPRE", "MPI", "BlockDiagonals", "Enzyme", "FiniteDiff", "BandedMatrices"]
+test = ["Test", "IterativeSolvers", "InteractiveUtils", "JET", "KrylovKit", "Pkg", "Random", "SafeTestsets", "MultiFloats", "ForwardDiff", "HYPRE", "MPI", "BlockDiagonals", "Enzyme", "EnzymeTestUtils", "FiniteDiff", "BandedMatrices"]

ext/LinearSolveEnzymeExt.jl

@@ -9,6 +9,47 @@ using Enzyme
 
 using EnzymeCore
 
+function EnzymeCore.EnzymeRules.forward(func::Const{typeof(LinearSolve.init)}, ::Type{RT}, prob::EnzymeCore.Annotation{LP}, alg::Const; kwargs...) where {RT, LP <: LinearSolve.LinearProblem}
+    @assert !(prob isa Const)
+    res = func.val(prob.val, alg.val; kwargs...)
+    if RT <: Const
+        return res
+    end
+    dres = func.val(prob.dval, alg.val; kwargs...)
+    dres.b .= res.b == dres.b ? zero(dres.b) : dres.b
+    dres.A .= res.A == dres.A ? zero(dres.A) : dres.A
+    if RT <: DuplicatedNoNeed
+        return dres
+    elseif RT <: Duplicated
+        return Duplicated(res, dres)
+    end
+end
+
+function EnzymeCore.EnzymeRules.forward(func::Const{typeof(LinearSolve.solve!)}, ::Type{RT}, linsolve::EnzymeCore.Annotation{LP}; kwargs...) where {RT, LP <: LinearSolve.LinearCache}
+    @assert !(linsolve isa Const)
+
+    A = deepcopy(linsolve.val.A) #mutates after function is applied
+    res = func.val(linsolve.val; kwargs...)
+    
+    if RT <: Const
+        return res
+    end
+    
+    dres = deepcopy(res)
+    invA = inv(A)
+    db = linsolve.dval.b
+    dA = linsolve.dval.A
+    dres.u .= invA * (db - dA * res.u)
+
+    if RT <: DuplicatedNoNeed
+        return dres
+    elseif RT <: Duplicated
+        return Duplicated(res, dres)
+    end
+
+    return Duplicated(res, dres)
+end
+
 function EnzymeCore.EnzymeRules.augmented_primal(config, func::Const{typeof(LinearSolve.init)}, ::Type{RT}, prob::EnzymeCore.Annotation{LP}, alg::Const; kwargs...) where {RT, LP <: LinearSolve.LinearProblem}
     res = func.val(prob.val, alg.val; kwargs...)
     dres = if EnzymeRules.width(config) == 1

test/enzyme.jl

@@ -1,5 +1,6 @@
 using Enzyme, ForwardDiff
 using LinearSolve, LinearAlgebra, Test
+using FiniteDiff
 
 n = 4
 A = rand(n, n);
@@ -161,4 +162,59 @@ Enzyme.autodiff(Reverse, f3, Duplicated(copy(A), dA), Duplicated(copy(b1), db1),
 @test dA ≈ dA2 atol=5e-5
 @test db1 ≈ db12
 @test db2 ≈ db22
-=#
+=#
+
+
+A = rand(n, n);
+dA = zeros(n, n);
+b1 = rand(n);
+function fb(b; alg = LUFactorization())
+    prob = LinearProblem(A, b)
+
+    sol1 = solve(prob, alg)
+
+    sum(sol1.u)
+end
+fb(b1)
+
+manual_jac = map(onehot(b1)) do db
+    y = A \ b1
+    sum(inv(A) * (db - dA*y))
+end |> collect
+@show manual_jac
+
+fd_jac = FiniteDiff.finite_difference_jacobian(fb, b1) |> vec
+@show fd_jac
+
+en_jac = map(onehot(b1)) do db1
+    eres = Enzyme.autodiff(Forward, fb, Duplicated(copy(b1), db1))
+    eres[1]
+end |> collect
+@show en_jac
+
+@test_broken en_jac ≈ manual_jac
+@test_broken en_jac ≈ fd_jac
+
+function fA(A; alg = LUFactorization())
+    prob = LinearProblem(A, b1)
+
+    sol1 = solve(prob, alg)
+
+    sum(sol1.u)
+end
+fA(A)
+
+manual_jac = map(onehot(A)) do dA
+    y = A \ b1
+    sum(inv(A) * (db1 - dA*y))
+end |> collect
+
+fd_jac = FiniteDiff.finite_difference_jacobian(fA, A) |> vec
+
+en_jac = map(onehot(A)) do dA
+    eres = Enzyme.autodiff(Forward, fA, Duplicated(copy(A), dA))
+    eres[1]
+end |> collect
+
+@test_broken en_jac ≈ manual_jac
+@test_broken en_jac ≈ fd_jac