Merge pull request #116 from DEEPDIP-project/Enzyme_2

Enzyme Implementation

CITATION.cff

@@ -9,6 +9,10 @@ authors:
     family-names: Agdestein
     affiliation: Centrum Wiskunde & Informatica
     orcid: 'https://orcid.org/0000-0002-1589-2916'
+  - given-names: Simone
+    family-names: Ciarella
+    affiliation: Netherlands eScience Center
+    orcid: 'https://orcid.org/0000-0002-9247-139X'
   - given-names: "Benjamin"
     family-names: "Sanderse"
     affiliation: Centrum Wiskunde & Informatica

Project.toml

@@ -7,6 +7,7 @@ version = "2.0.1"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
+EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 IterativeSolvers = "42fd0dbc-a981-5370-80f2-aaf504508153"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
@@ -36,6 +37,7 @@ CUDA = "5"
 CUDSS = "0.3"
 ChainRulesCore = "1"
 DocStringExtensions = "0.9"
+EnzymeCore = "0.8"
 FFTW = "1"
 IterativeSolvers = "0.9"
 KernelAbstractions = "0.9"

docs/Project.toml

@@ -2,16 +2,12 @@
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 DocumenterCitations = "daee34ce-89f3-4625-b898-19384cb65244"
 DocumenterVitepress = "4710194d-e776-4893-9690-8d956a29c365"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 Examples = "318dbb63-4243-420f-99f2-d56058123f9d"
 IncompressibleNavierStokes = "5e318141-6589-402b-868d-77d7df8c442e"
 Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
 NeuralClosure = "099dac27-d7f2-4047-93d5-0baee36b9c25"
 
-[sources]
-Examples = {path = "../examples"}
-IncompressibleNavierStokes = {path = ".."}
-NeuralClosure = {path = "../lib/NeuralClosure"}
-
 [compat]
 Documenter = "1"
 DocumenterCitations = "1"
@@ -21,3 +17,12 @@ IncompressibleNavierStokes = "2"
 Literate = "2"
 NeuralClosure = "1"
 julia = "1.9"
+
+[sources.Examples]
+path = "../examples"
+
+[sources.IncompressibleNavierStokes]
+path = ".."
+
+[sources.NeuralClosure]
+path = "../lib/NeuralClosure"

docs/src/manual/differentiability.md

@@ -7,19 +7,24 @@ CurrentModule = IncompressibleNavierStokes
 IncompressibleNavierStokes is
 [reverse-mode differentiable](https://juliadiff.org/ChainRulesCore.jl/stable/index.html#Reverse-mode-AD-rules-(rrules)),
 which means that you can back-propagate gradients through the code.
+Two AD libraries are currently supported:
+* **[Zygote.jl](https://github.com/FluxML/Zygote.jl)**: it is the default AD library in the Julia ecosystem and is the most widely used.
+* **[Enzyme.jl](https://github.com/EnzymeAD/Enzyme.jl)**: currently has low coverage over the Julia programming language, however it is usually the most efficient if applicable.
+
+## Automatic differentiation with Zygote
+
+Zygote.jl is the default choice for AD backend because it is easy to understand, compatible with most of the Julia ecosystem and good with vectorized code and BLAS.
 This comes at a cost however, as intermediate velocity fields need to be stored
 in memory for use in the backward pass. For this reason, many of the operators
 come in two versions: a slow differentiable allocating non-mutating variant (e.g.
 [`divergence`](@ref)) and fast non-differentiable non-allocating mutating
 variant (e.g. [`divergence!`](@ref).)
 
-!!! warning "Differentiable code"
+!!! warning "Zygote limitation: array mutation"
     To make your code differentiable, you must use the differentiable versions
     of the operators (without the exclamation marks).
 
-To differentiate the code, use [Zygote.jl](https://github.com/FluxML/Zygote.jl).
-
-## Example: Gradient of kinetic energy
+#### Example: Gradient of kinetic energy
 
 To differentiate outputs of a simulation with respect to the initial conditions,
 make a time stepping loop composed of differentiable operations:
@@ -55,3 +60,56 @@ Now `g` is the gradient of `final_energy` with respect to the initial conditions
 
 Note that every operation in the `final_energy` function is non-mutating and
 thus differentiable.
+
+--- 
+## Automatic differentiation with Enzyme
+
+Enzyme.jl is highly-efficient and its ability to perform AD on optimized code allows Enzyme to meet or exceed the performance of state-of-the-art AD tools.
+The downside is that restricts the user's defined f function to not do things like require garbage collection or calls to BLAS/LAPACK. However, mutation is supported, meaning that in-place f with fully mutating non-allocating code will work with Enzyme and this will be the most efficient adjoint implementation.
+
+!!! warning "Enzyme limitation: vector returns"
+    Enzyme's autodiff function can only handle functions with scalar output. To implement pullbacks for array-valued functions, use a mutating function that returns `nothing` and stores its result in one of the arguments, which must be passed wrapped in a Duplicated.
+    In IncompressibleNavierStokes, we provide `enzyme_wrapper` to automatically wrap the function and its arguments in the correct way.
+
+#### Example: Gradient of the right-hand side
+
+In this example we differentiate the right-hand side of the Navier-Stokes equations with respect to the velocity field `u`:
+
+```@example
+import IncompressibleNavierStokes as INS
+using Enzyme
+ax = range(0, 1, 101)
+setup = INS.Setup(; x = (ax, ax), Re = 500.0)
+psolver = INS.default_psolver(setup)
+u = INS.random_field(setup)
+dudt = similar(u)
+t = 0.0
+f! = INS.right_hand_side!
+```
+Notice that we are using the mutating (in-place) version of the right-hand side function. This function can not be differentiate by Zygote, which requires the slower non-mutating version of the right-hand side.
+
+We then define the `Dual` part of the input and output, required to store the adjoint values: 
+```@example
+ddudt = Enzyme.make_zero(dudt) .+ 1;
+du = Enzyme.make_zero(u);
+```
+Remember that the derivative of the output (also called the *seed*) has to be set to $1$ in order to compute the gradient. In this case the output is the force, that we store mutating the value of `dudt` inside `right_hand_side!`.
+
+Then we pack the parameters to be passed to `right_hand_side!`:
+```@example
+params = [setup, psolver];
+params_ref = Ref(params);
+```
+Now, we call the `autodiff` function from Enzyme:
+```@example
+Enzyme.autodiff(Enzyme.Reverse, f!, Duplicated(dudt,ddudt), Duplicated(u,du), Const(params_ref), Const(t))
+```
+Since we have passed a `Duplicated` object, the gradient of `u` is stored in `du`. 
+
+Finally, we can also compare its value with the one obtained by Zygote differentiating the out-of-place (non-mutating) version of the right-hand side:
+```@example
+using Zygote
+f = create_right_hand_side(setup, psolver)
+_, zpull = Zygote.pullback(f, u, nothing, T(0));
+@assert zpull(dudt)[1] == du
+```

docs/src/manual/sciml.md

@@ -0,0 +1,46 @@
+```@meta
+CurrentModule = IncompressibleNavierStokes
+```
+
+# Using IncompressibleNavierStokes in SciML
+
+The [SciML organization](https://sciml.ai/) is a collection of tools for solving equations and modeling systems. It has a coherent development principle, unified APIs over large collections of equation solvers, pervasive differentiability and sensitivity analysis, and features many of the highest performance and parallel implementations one can find.
+
+In particular, [DifferentialEquations.jl](https://docs.sciml.ai/DiffEqDocs/stable/) contains tools to solve differential equations defined as $\dfrac{du}{dt} = f(u, t)$ that include a large collection of solvers, sensitivity analysis, and more.
+
+Using IncompressibleNavierStokes it is possible to write the momentum equations without the pressure by explicitly solving the discrete Poisson equation and obtaining:
+
+```math
+\begin{align*}
+\frac{\mathrm{d} u_h}{\mathrm{d} t} &= (I - G L^{-1} W M)
+(F(u_h) - y_G) - G L^{-1} W \frac{\mathrm{d} y_M}{\mathrm{d} t}\\ &=f(u_h).
+\end{align*}
+```
+
+The derivation and the drawbacks of this approach are discussed in the [documentation](/docs/src/manual/spatial.md).
+
+This projected right-hand side can be used in the SciML solvers to solve the Navier-Stokes equations. The following example shows how to use the SciML solvers to solve the ODEs obtained from the Navier-Stokes equations.
+
+```julia
+using DifferentialEquations 
+f(u, p, t) = create_right_hand_side(setup, psolver) 
+u0 = INITIAL_CONDITION
+tspan = (0.0, 1.0)     # time span where to solve.
+problem = ODEProblem(f, u0, tspan) #SciMLBase.ODEProblem
+sol = solve(problem, Tsit5(), reltol = 1e-8, abstol = 1e-8) # sol: SciMLBase.ODESolution
+```
+
+Alternatively, it is also possible to use an [in-place formulation](https://docs.sciml.ai/DiffEqDocs/stable/basics/problem/#In-place-vs-Out-of-Place-Function-Definition-Forms) 
+
+```julia
+f(du,u,p,t) = right_hand_side!(du, u, Ref([setup, psolver]), t)
+```
+that is usually faster than the out-of-place formulation.
+
+You can look [here](https://docs.sciml.ai/DiffEqDocs/stable/basics/overview/) for more information on how to use the SciML solvers and all the options available.
+
+## API 
+```@autodocs
+Modules = [IncompressibleNavierStokes]
+Pages = ["sciml.jl"]
+```

src/IncompressibleNavierStokes.jl

@@ -13,6 +13,8 @@ using Adapt
 using ChainRulesCore
 using DocStringExtensions
 using FFTW
+using EnzymeCore
+using EnzymeCore.EnzymeRules
 using IterativeSolvers
 using KernelAbstractions
 using KernelAbstractions.Extras.LoopInfo: @unroll
@@ -66,6 +68,7 @@ include("eddyviscosity.jl")
 include("matrices.jl")
 include("initializers.jl")
 include("processors.jl")
+include("sciml.jl")
 include("solver.jl")
 include("utils.jl")
 
@@ -119,6 +122,7 @@ export apply_bc_u,
     apply_bc_p,
     apply_bc_temp,
     applybodyforce,
+    applypressure,
     convection_diffusion_temp,
     convection,
     diffusion,
@@ -135,6 +139,7 @@ export apply_bc_u,
     laplacian_mat,
     momentum,
     poisson,
+    pressure,
     pressuregradient,
     project,
     scalewithvolume,
@@ -145,4 +150,7 @@ export apply_bc_u,
     Dfield,
     Qfield
 
+# SciML operations
+export create_right_hand_side, right_hand_side!
+
 end

src/boundary_conditions.jl

@@ -155,7 +155,6 @@ ChainRulesCore.rrule(::typeof(apply_bc_temp), temp, t, setup) = (
         NoTangent(),
     ),
 )
-
 "Apply velocity boundary conditions (in-place version)."
 function apply_bc_u!(u, t, setup; kwargs...)
     (; boundary_conditions) = setup
@@ -528,3 +527,80 @@ apply_bc_temp!(bc::PressureBC, temp, β, t, setup; isright, kwargs...) =
 
 apply_bc_temp_pullback!(bc::PressureBC, φbar, β, t, setup; isright, kwargs...) =
     apply_bc_p_pullback!(SymmetricBC(), φbar, β, t, setup; isright, kwargs...)
+
+# COV_EXCL_START
+# Wrap a function to return `nothing`, because Enzyme can not handle vector return values.
+function enzyme_wrap(
+    f::Union{typeof(apply_bc_u!),typeof(apply_bc_p!),typeof(apply_bc_temp!)},
+)
+    # the boundary condition modifies x which is usually the field that we want to differentiate, so we need to introduce a copy of it and modify it instead
+    function wrapped_f(y, x, args...)
+        y .= x
+        f(y, args...)
+        return nothing
+    end
+    return wrapped_f
+end
+
+function EnzymeRules.augmented_primal(
+    config::RevConfigWidth{1},
+    func::Union{
+        Const{typeof(enzyme_wrap(apply_bc_u!))},
+        Const{typeof(enzyme_wrap(apply_bc_p!))},
+        Const{typeof(enzyme_wrap(apply_bc_temp!))},
+    },
+    ::Type{<:Const},
+    y::Duplicated,
+    x::Duplicated,
+    t::Const,
+    setup::Const,
+)
+    primal = func.val(y.val, x.val, t.val, setup.val)
+    return AugmentedReturn(primal, nothing, nothing)
+end
+function EnzymeRules.reverse(
+    config::RevConfigWidth{1},
+    func::Const{typeof(enzyme_wrap(apply_bc_u!))},
+    dret,
+    tape,
+    y::Duplicated,
+    x::Duplicated,
+    t::Const,
+    setup::Const,
+)
+    adj = apply_bc_u_pullback!(x.val, t.val, setup.val)
+    x.dval .+= adj
+    y.dval .= x.dval # y is a copy of x
+    return (nothing, nothing, nothing, nothing)
+end
+function EnzymeRules.reverse(
+    config::RevConfigWidth{1},
+    func::Const{typeof(enzyme_wrap(apply_bc_p!))},
+    dret,
+    tape,
+    y::Duplicated,
+    x::Duplicated,
+    t::Const,
+    setup::Const,
+)
+    adj = apply_bc_p_pullback!(x.val, t.val, setup.val)
+    x.dval .+= adj
+    y.dval .= x.dval # y is a copy of x
+    return (nothing, nothing, nothing, nothing)
+end
+function EnzymeRules.reverse(
+    config::RevConfigWidth{1},
+    func::Const{typeof(enzyme_wrap(apply_bc_temp!))},
+    dret,
+    tape,
+    y::Duplicated,
+    x::Duplicated,
+    t::Const,
+    setup::Const,
+)
+    adj = apply_bc_temp_pullback!(x.val, t.val, setup.val)
+    x.dval .+= adj
+    y.dval .= x.dval # y is a copy of x
+    return (nothing, nothing, nothing, nothing)
+end
+# COV_EXCL_STOP