Merge pull request #2742 from SciML/mooncake_direct

Mooncake Direct Adjoints for Explicit Integrators

Author: ChrisRackauckas

lib/OrdinaryDiffEqCore/Project.toml

@@ -39,9 +39,11 @@ TruncatedStacktraces = "781d530d-4396-4725-bb49-402e4bee1e77"
 
 [weakdeps]
 EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 
 [extensions]
 OrdinaryDiffEqCoreEnzymeCoreExt = "EnzymeCore"
+OrdinaryDiffEqCoreMooncakeExt = "Mooncake"
 
 [compat]
 ADTypes = "0.2, 1"
@@ -63,6 +65,7 @@ InteractiveUtils = "1.9"
 LinearAlgebra = "1.9"
 Logging = "1.9"
 MacroTools = "0.5"
+Mooncake = "0.4"
 MuladdMacro = "0.2.1"
 Polyester = "0.7"
 PrecompileTools = "1"

lib/OrdinaryDiffEqCore/ext/OrdinaryDiffEqCoreEnzymeCoreExt.jl

@@ -26,4 +26,9 @@ function EnzymeCore.EnzymeRules.inactive_noinl(
     true
 end
 
+function EnzymeCore.EnzymeRules.inactive_noinl(
+        ::typeof(OrdinaryDiffEqCore.final_progress), args...)
+    true
+end
+
 end

lib/OrdinaryDiffEqCore/ext/OrdinaryDiffEqCoreMooncakeExt.jl

@@ -0,0 +1,10 @@
+module OrdinaryDiffEqCoreMooncakeExt
+
+using OrdinaryDiffEqCore, Mooncake
+using Mooncake: @zero_adjoint, MinimalCtx
+Mooncake.@zero_adjoint Mooncake.MinimalCtx Tuple{typeof(OrdinaryDiffEqCore.ode_determine_initdt), Any, Any, Any, Any, Any, Any, Any, Any, Any}
+Mooncake.@zero_adjoint Mooncake.MinimalCtx Tuple{typeof(OrdinaryDiffEqCore.SciMLBase.check_error), Any}
+Mooncake.@zero_adjoint Mooncake.MinimalCtx Tuple{typeof(OrdinaryDiffEqCore.fixed_t_for_floatingpoint_error!), Any, Any}
+Mooncake.@zero_adjoint Mooncake.MinimalCtx Tuple{typeof(OrdinaryDiffEqCore.final_progress), Any}
+
+end

lib/OrdinaryDiffEqCore/src/integrators/integrator_utils.jl

@@ -195,15 +195,19 @@ function _postamble!(integrator)
         resize!(integrator.sol.k, integrator.saveiter_dense)
     end
     if integrator.opts.progress
-        @logmsg(LogLevel(-1),
-            integrator.opts.progress_name,
-            _id=integrator.opts.progress_id,
-            message=integrator.opts.progress_message(integrator.dt, integrator.u,
-                integrator.p, integrator.t),
-            progress="done")
+
     end
 end
 
+function final_progress(integrator)
+    @logmsg(LogLevel(-1),
+    integrator.opts.progress_name,
+    _id=integrator.opts.progress_id,
+    message=integrator.opts.progress_message(integrator.dt, integrator.u,
+        integrator.p, integrator.t),
+    progress="done")
+end
+
 function solution_endpoint_match_cur_integrator!(integrator)
     if integrator.opts.save_end &&
        (integrator.saveiter == 0 ||

test/downstream/Project.toml

@@ -3,6 +3,7 @@ DDEProblemLibrary = "f42792ee-6ffc-4e2a-ae83-8ee2f22de800"
 DelayDiffEq = "bcd4f6db-9728-5f36-b5f7-82caef46ccdb"
 Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 Measurements = "eff96d63-e80a-5855-80a2-b1b0885c5ab7"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 StochasticDiffEq = "789caeaf-c7a9-5a7d-9973-96adeb23e2a0"

test/downstream/mooncake.jl

@@ -0,0 +1,19 @@
+using Mooncake, OrdinaryDiffEq, StaticArrays
+
+function lorenz!(du, u, p, t)
+    du[1] = 10.0(u[2] - u[1])
+    du[2] = u[1] * (28.0 - u[3]) - u[2]
+    du[3] = u[1] * u[2] - (8 / 3) * u[3]
+end
+
+const _saveat =  SA[0.0,0.25,0.5,0.75,1.0,1.25,1.5,1.75,2.0,2.25,2.5,2.75,3.0]
+
+function f(u0::Array{Float64})
+    tspan = (0.0, 3.0)
+    prob = ODEProblem{true, SciMLBase.FullSpecialize}(lorenz!, u0, tspan)
+    sol = DiffEqBase.solve(prob, Tsit5(), saveat = _saveat, sensealg = DiffEqBase.SensitivityADPassThrough())
+    sum(sol)
+end;
+u0 = [1.0; 0.0; 0.0]
+mooncake_gradient(f, x) = Mooncake.value_and_gradient!!(Mooncake.build_rrule(f, x), f, x)[2][2]
+mooncake_gradient(f, u0)

test/runtests.jl

@@ -156,6 +156,9 @@ end
         activate_downstream_env()
         @time @safetestset "DelayDiffEq Tests" include("downstream/delaydiffeq.jl")
         @time @safetestset "Measurements Tests" include("downstream/measurements.jl")
+        if VERSION >= v"1.11"
+            @time @safetestset "Mooncake Tests" include("downstream/mooncake.jl")
+        end
         @time @safetestset "Sparse Diff Tests" include("downstream/sparsediff_tests.jl")
         @time @safetestset "Time derivative Tests" include("downstream/time_derivative_test.jl")
     end