Merge pull request #291 from ErikQQY/qqy/more_AD

Support Enzyme and Mooncake in AD backends

Author: ChrisRackauckas

docs/src/basics/autodiff.md

@@ -2,7 +2,7 @@
 
 !!! note
 
-    We support ForwardDiff.jl, FiniteDiff.jl and PolyesterForwardDiff.jl(PolyesterForwardDiff only for collocation methods) via DifferentiationInterface.jl. Please refer to
+    We support ForwardDiff.jl, FiniteDiff.jl, Enzyme.jl, Mooncake.jl and PolyesterForwardDiff.jl(PolyesterForwardDiff only for collocation methods) via DifferentiationInterface.jl. Please refer to
     the [backends page](https://juliadiff.org/DifferentiationInterface.jl/DifferentiationInterface/stable/explanation/backends/)
     for more information.
 
@@ -25,3 +25,9 @@ In BoundaryValueDiffEq.jl, we require AD to obtain the Jacobian of the loss func
   - [`AutoPolyesterForwardDiff`](@extref ADTypes.AutoPolyesterForwardDiff): Might be faster
     than [`AutoForwardDiff`](@extref ADTypes.AutoForwardDiff) for large problems. Requires
     `PolyesterForwardDiff.jl` to be installed and loaded.
+  - [`AutoEnzyme(; mode = Enzyme.Forward)`](@extref ADTypes.AutoEnzyme): Source transformation forward mode AD.
+
+## Summary of Reverse Mode AD Backends
+
+  - [`AutoEnzyme(; mode = Enzyme.Reverse)`](@extref ADTypes.AutoEnzyme): Source transformation reverse mode AD.
+  - [`AutoMooncake(; config = nothing)`](@extref ADTypes.AutoMooncake): Source transformation reverse mode AD.

docs/src/devdocs/internal_interfaces.md

@@ -2,6 +2,6 @@
 
 ## Solvers
 
-```julia
+```@docs
 BoundaryValueDiffEqCore.AbstractBoundaryValueDiffEqAlgorithm
 ```

lib/BoundaryValueDiffEqAscher/src/BoundaryValueDiffEqAscher.jl

@@ -2,13 +2,15 @@ module BoundaryValueDiffEqAscher
 
 using ADTypes: ADTypes, AutoSparse, AutoForwardDiff
 using AlmostBlockDiagonals: AlmostBlockDiagonals, IntermediateAlmostBlockDiagonal
+
 using BoundaryValueDiffEqCore: AbstractBoundaryValueDiffEqAlgorithm,
                                AbstractBoundaryValueDiffEqCache, BVPJacobianAlgorithm,
                                __extract_problem_details, concrete_jacobian_algorithm,
                                __Fix3, __concrete_nonlinearsolve_algorithm,
                                __internal_nlsolve_problem, __vec, __vec_f, __vec_f!,
                                __vec_bc, __vec_bc!, __extract_mesh, get_dense_ad,
-                               __get_bcresid_prototype, __split_kwargs
+                               __get_bcresid_prototype, __split_kwargs,
+                               __default_nonsparse_ad
 
 using ConcreteStructs: @concrete
 using DiffEqBase: DiffEqBase

lib/BoundaryValueDiffEqAscher/src/algorithms.jl

@@ -55,7 +55,7 @@ for stage in (1, 2, 3, 4, 5, 6, 7)
     end
 end
 
-function concretize_jacobian_algorithm(alg::AbstractAscher, prob)
-    @set! alg.jac_alg = concrete_jacobian_algorithm(alg.jac_alg, prob, alg)
-    return alg
+function BoundaryValueDiffEqCore.concrete_jacobian_algorithm(
+        jac_alg::BVPJacobianAlgorithm, prob::BVProblem, alg::AbstractAscher)
+    return BVPJacobianAlgorithm(__default_nonsparse_ad(prob.u0))
 end

lib/BoundaryValueDiffEqCore/src/BoundaryValueDiffEqCore.jl

@@ -18,7 +18,7 @@ using SciMLBase: SciMLBase, AbstractBVProblem, AbstractDiffEqInterpolation,
                  StandardBVProblem, StandardSecondOrderBVProblem, __solve, _unwrap_val
 using Setfield: @set!, @set
 using SparseArrays: sparse
-using SparseConnectivityTracer: TracerLocalSparsityDetector
+using SparseConnectivityTracer: SparseConnectivityTracer, TracerLocalSparsityDetector
 using SparseMatrixColorings: GreedyColoringAlgorithm
 
 @reexport using NonlinearSolveFirstOrder, SciMLBase

lib/BoundaryValueDiffEqCore/src/types.jl

@@ -68,15 +68,25 @@ function concrete_jacobian_algorithm(
     return concrete_jacobian_algorithm(jac_alg, prob.problem_type, prob, alg)
 end
 
+# For multi-point BVP, we only care about bc_diffmode and nonbc_diffmode
 function concrete_jacobian_algorithm(
-        jac_alg::BVPJacobianAlgorithm, prob_type, prob::BVProblem, alg)
+        jac_alg::BVPJacobianAlgorithm, prob_type::StandardBVProblem, prob::BVProblem, alg)
     u0 = __extract_u0(prob.u0, prob.p, first(prob.tspan))
-    diffmode = jac_alg.diffmode === nothing ? __default_sparse_ad(u0) : jac_alg.diffmode
-    bc_diffmode = jac_alg.bc_diffmode === nothing ?
-                  (prob_type isa TwoPointBVProblem ? __default_bc_sparse_ad :
-                   __default_nonsparse_ad)(u0) : jac_alg.bc_diffmode
+    bc_diffmode = jac_alg.bc_diffmode === nothing ? __default_bc_sparse_ad(u0) :
+                  jac_alg.bc_diffmode
     nonbc_diffmode = jac_alg.nonbc_diffmode === nothing ? __default_sparse_ad(u0) :
                      jac_alg.nonbc_diffmode
+    diffmode = jac_alg.diffmode === nothing ? nothing : jac_alg.diffmode
+    return BVPJacobianAlgorithm(bc_diffmode, nonbc_diffmode, diffmode)
+end
+
+# For two-point BVP, we only care about diffmode
+function concrete_jacobian_algorithm(
+        jac_alg::BVPJacobianAlgorithm, prob_type::TwoPointBVProblem, prob::BVProblem, alg)
+    u0 = __extract_u0(prob.u0, prob.p, first(prob.tspan))
+    diffmode = jac_alg.diffmode === nothing ? __default_sparse_ad(u0) : jac_alg.diffmode
+    bc_diffmode = jac_alg.bc_diffmode === nothing ? nothing : jac_alg.bc_diffmode
+    nonbc_diffmode = jac_alg.nonbc_diffmode === nothing ? nothing : jac_alg.nonbc_diffmode
     return BVPJacobianAlgorithm(bc_diffmode, nonbc_diffmode, diffmode)
 end
 
@@ -151,22 +161,13 @@ end
            __needs_diffcache(jac_alg.nonbc_diffmode)
 end
 
-# We don't need to always allocate a DiffCache. This works around that.
-@concrete struct FakeDiffCache
-    du
-end
-
 function __maybe_allocate_diffcache(x, chunksize, jac_alg)
-    return __needs_diffcache(jac_alg) ? DiffCache(x, chunksize) : FakeDiffCache(x)
+    return __needs_diffcache(jac_alg) ? DiffCache(x, chunksize) : x
 end
 __maybe_allocate_diffcache(x::DiffCache, chunksize) = DiffCache(zero(x.du), chunksize)
-__maybe_allocate_diffcache(x::FakeDiffCache, _) = FakeDiffCache(zero(x.du))
-
-const MaybeDiffCache = Union{DiffCache, FakeDiffCache}
 
 ## get_tmp shows a warning as it should on cache exapansion, this behavior however is
 ## expected for adaptive BVP solvers so we write our own `get_tmp` and drop the warning logs
-@inline get_tmp(dc::FakeDiffCache, u) = dc.du
 
 @inline function get_tmp(dc, u)
     return Logging.with_logger(Logging.NullLogger()) do
@@ -180,4 +181,8 @@ struct NoDiffCacheNeeded end
 
 @inline __cache_trait(::AutoForwardDiff) = DiffCacheNeeded()
 @inline __cache_trait(ad::AutoSparse) = __cache_trait(ADTypes.dense_ad(ad))
+@inline function __cache_trait(jac_alg::BVPJacobianAlgorithm)
+    isnothing(jac_alg.diffmode) ? __cache_trait(jac_alg.nonbc_diffmode) :
+    __cache_trait(jac_alg.diffmode)
+end
 @inline __cache_trait(_) = NoDiffCacheNeeded()

lib/BoundaryValueDiffEqCore/src/utils.jl

@@ -1,5 +1,5 @@
 recursive_length(x::Vector{<:AbstractArray}) = sum(length, x)
-recursive_length(x::Vector{<:MaybeDiffCache}) = sum(xᵢ -> length(xᵢ.du), x)
+recursive_length(x::Vector{<:DiffCache}) = sum(xᵢ -> length(xᵢ.u), x)
 
 function recursive_flatten(x::Vector{<:AbstractArray})
     y = zero(first(x), recursive_length(x))
@@ -37,7 +37,7 @@ end
     return y
 end
 
-@views function recursive_unflatten!(y::Vector{<:MaybeDiffCache}, x::AbstractVector)
+@views function recursive_unflatten!(y::Vector{<:DiffCache}, x::AbstractVector)
     return recursive_unflatten!(get_tmp.(y, (x,)), x)
 end
 
@@ -230,7 +230,7 @@ end
 __resize!(::Nothing, n, _) = nothing
 __resize!(::Nothing, n, _, _) = nothing
 
-function __resize!(x::AbstractVector{<:MaybeDiffCache}, n, M)
+function __resize!(x::AbstractVector{<:DiffCache}, n, M)
     N = n - length(x)
     N == 0 && return x
     if N > 0
@@ -567,6 +567,14 @@ function _sparse_like(I, J, x::AbstractArray, m = maximum(I), n = maximum(J))
     return sparse(I′, J′, V, m, n)
 end
 
+nodual_value(x) = x
+nodual_value(x::ForwardDiff.Dual) = ForwardDiff.value(x)
+nodual_value(x::AbstractArray{<:ForwardDiff.Dual}) = map(ForwardDiff.value, x)
+nodual_value(x::SparseConnectivityTracer.Dual) = SparseConnectivityTracer.primal(x)
+function nodual_value(x::AbstractArray{<:SparseConnectivityTracer.Dual})
+    map(SparseConnectivityTracer.primal, x)
+end
+
 function __split_kwargs(; abstol, adaptive, controller, kwargs...)
     return ((abstol, adaptive, controller), (; abstol, adaptive, kwargs...))
 end

lib/BoundaryValueDiffEqFIRK/Project.toml

@@ -36,6 +36,7 @@ ConcreteStructs = "0.2.3"
 DiffEqBase = "6.167"
 DiffEqDevTools = "2.44"
 DifferentiationInterface = "0.6.42"
+Enzyme = "0.13.33"
 FastAlmostBandedMatrices = "0.1.4"
 FastClosures = "0.3.2"
 ForwardDiff = "0.10.38, 1"
@@ -44,6 +45,7 @@ InteractiveUtils = "<0.0.1, 1"
 JET = "0.9.18"
 LinearAlgebra = "1.10"
 LinearSolve = "2.36.2, 3"
+Mooncake = "0.4.108"
 OrdinaryDiffEqRosenbrock = "1"
 PreallocationTools = "0.4.24"
 PrecompileTools = "1.2"
@@ -62,10 +64,12 @@ julia = "1.10"
 [extras]
 Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595"
 DiffEqDevTools = "f3b72e0c-5b89-59e1-b016-84e28bfd966d"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 Hwloc = "0e44f5e4-bd66-52a0-8798-143a42290a1d"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 JET = "c3a54625-cd67-489e-a8e7-0a5a0ff4e31b"
 LinearSolve = "7ed4a6bd-45f5-4d41-b270-4a48e9bafcae"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 OrdinaryDiffEqRosenbrock = "43230ef6-c299-4910-a778-202eb28ce4ce"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReTestItems = "817f1d60-ba6b-4fd5-9520-3cf149f6a823"
@@ -74,4 +78,4 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Aqua", "DiffEqDevTools", "Hwloc", "InteractiveUtils", "JET", "LinearSolve", "OrdinaryDiffEqRosenbrock", "Random", "ReTestItems", "RecursiveArrayTools", "StaticArrays", "Test"]
+test = ["Aqua", "DiffEqDevTools", "Enzyme", "Hwloc", "InteractiveUtils", "JET", "LinearSolve", "OrdinaryDiffEqRosenbrock", "Mooncake", "Random", "ReTestItems", "RecursiveArrayTools", "StaticArrays", "Test"]

lib/BoundaryValueDiffEqFIRK/src/BoundaryValueDiffEqFIRK.jl

@@ -9,16 +9,20 @@ using BoundaryValueDiffEqCore: AbstractBoundaryValueDiffEqAlgorithm,
                                __concrete_nonlinearsolve_algorithm, diff!, EvalSol,
                                concrete_jacobian_algorithm, eval_bc_residual, interval,
                                eval_bc_residual!, get_tmp, __maybe_matmul!, __resize!,
-                               __extract_problem_details, __initial_guess,
+                               __extract_problem_details, __initial_guess, nodual_value,
+                               __maybe_allocate_diffcache, __restructure_sol,
+                               __get_bcresid_prototype, __vec, __vec_f, __vec_f!, __vec_bc,
+                               __vec_bc!, recursive_flatten_twopoint!,
+                               __internal_nlsolve_problem, __extract_mesh, __extract_u0,
                                __default_coloring_algorithm, __maybe_allocate_diffcache,
                                __restructure_sol, __get_bcresid_prototype, safe_similar,
-                               __vec, __vec_f, __vec_f!, __vec_bc, __vec_bc!,
+                               __vec, __vec_f, __vec_f!, __vec_bc, __vec_bc!, __cache_trait,
                                recursive_flatten_twopoint!, __internal_nlsolve_problem,
-                               MaybeDiffCache, __extract_mesh, __extract_u0,
-                               __has_initial_guess, __initial_guess_length,
-                               __initial_guess_on_mesh, __flatten_initial_guess,
-                               __split_kwargs, __build_solution, __Fix3, _sparse_like,
-                               get_dense_ad
+                               __extract_mesh, __extract_u0, DiffCacheNeeded,
+                               NoDiffCacheNeeded, __has_initial_guess,
+                               __initial_guess_length, __initial_guess_on_mesh,
+                               __flatten_initial_guess, __build_solution, __Fix3,
+                               __split_kwargs, _sparse_like, get_dense_ad
 
 using ConcreteStructs: @concrete
 using DiffEqBase: DiffEqBase