Merge pull request #927 from wsmoses/master

Add Enzyme extension

Author: ChrisRackauckas

Project.toml

@@ -10,6 +10,7 @@ DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 EnumX = "4e289a0a-7415-4d19-859d-a7e5c4648b56"
+EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 FastBroadcast = "7034ab61-46d4-4ed7-9d0f-46aef9175898"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 FunctionWrappers = "069b7b12-0de2-55c6-9aab-29f3d0a68a2e"
@@ -38,6 +39,7 @@ ZygoteRules = "700de1a5-db45-46bc-99cf-38207098b444"
 
 [weakdeps]
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 GeneralizedGenerated = "6b9d7cbe-bcb9-11e9-073f-15a7a543e2eb"
 MPI = "da04e1cc-30fd-572f-bb4f-1f8673147195"
 Measurements = "eff96d63-e80a-5855-80a2-b1b0885c5ab7"
@@ -57,6 +59,7 @@ DiffEqBaseReverseDiffExt = "ReverseDiff"
 DiffEqBaseTrackerExt = "Tracker"
 DiffEqBaseUnitfulExt = "Unitful"
 DiffEqBaseZygoteExt = "Zygote"
+DiffEqBaseEnzymeExt = "Enzyme"
 
 [compat]
 ArrayInterface = "7"

ext/DiffEqBaseEnzymeExt.jl

@@ -0,0 +1,53 @@
+module DiffEqBaseEnzymeExt
+
+using DiffEqBase
+import DiffEqBase: value
+isdefined(Base, :get_extension) ? (import Enzyme) : (import ..Enzyme)
+
+using ChainRulesCore
+using EnzymeCore
+
+function EnzymeCore.EnzymeRules.augmented_primal(config::EnzymeCore.EnzymeRules.ConfigWidth{1}, func::Const{typeof(DiffEqBase.solve_up)}, ::Type{Duplicated{RT}}, prob, sensealg::Union{Const{Nothing}, Const{<:DiffEqBase.AbstractSensitivityAlgorithm}}, u0, p, args...; kwargs...) where RT
+    @inline function copy_or_reuse(val, idx)
+        if EnzymeCore.EnzymeRules.overwritten(config)[idx] && ismutable(val)
+            return deepcopy(val)
+        else
+            return val
+        end
+    end
+
+    @inline function arg_copy(i)
+        copy_or_reuse(args[i].val, i+5)
+    end
+ 
+    res = DiffEqBase._solve_adjoint(copy_or_reuse(prob.val, 2), copy_or_reuse(sensealg.val, 3), copy_or_reuse(u0.val, 4), copy_or_reuse(p.val, 5), SciMLBase.ChainRulesOriginator(), ntuple(arg_copy, Val(length(args)))...;
+        kwargs...)
+
+    dres = deepcopy(res[1])::RT
+    for v in dres.u
+        v.= 0
+    end
+    tup = (dres, res[2])
+    return EnzymeCore.EnzymeRules.AugmentedReturn{RT, RT, Any}(res[1], dres, tup::Any)
+end
+
+function EnzymeCore.EnzymeRules.reverse(config::EnzymeCore.EnzymeRules.ConfigWidth{1}, func::Const{typeof(DiffEqBase.solve_up)}, ::Type{<:Duplicated{RT}}, tape, prob, sensealg::Union{Const{Nothing}, Const{<:DiffEqBase.AbstractSensitivityAlgorithm}}, u0, p, args...; kwargs...) where RT
+	dres, clos = tape
+    dres = dres::RT
+	dargs = clos(dres)
+    for (darg, ptr) in zip(dargs, (func, prob, sensealg, u0, p, args...))
+        if ptr isa EnzymeCore.Const
+            continue
+        end
+        if darg == ChainRulesCore.NoTangent()
+            continue
+        end
+        ptr.dval .+= darg
+    end
+    for v in dres.u
+        v.= 0
+    end
+    return ntuple(_ -> nothing, Val(length(args)+4))
+end
+
+end

src/init.jl

@@ -47,5 +47,9 @@ end
         @require MPI="da04e1cc-30fd-572f-bb4f-1f8673147195" begin
             include("../ext/DiffEqBaseMPIExt.jl")
         end
+        
+        @require Enzyme="7da242da-08ed-463a-9acd-ee780be4f1d9" begin
+            include("../ext/DiffEqBaseEnzymeExt.jl")
+        end
     end
 end