The TrixiEnzyme Module

TrixiEnzyme

TrixiEnzyme.jl is a component package of the Trixi.jl ecosystem and integrates Trixi.jl with Compiler-Based (LLVM level) automatic differentiation via Enzyme.jl for hyperbolic partial differential equations (PDEs). The integration of Trixi.jl with Compiler-Based (LLVM level) automatic differentiation via Enzyme.jl offers the following benefits: facilitates rapid forward mode AD, enables reverse mode AD, supports cross-language AD, and critically, supports mutating operations and caching, on which Trixi.jl relies, to enhance the performance of both simulation runs and AD. The final deliverable will include as many of Trixi's advanced features as possible, such as adaptive mesh refinement, shock capturing, etc., showcasing the benefits of differentiable programming in Julia's ecosystem.

enzyme_rhs!(du_ode::AbstractVector, u_ode::AbstractVector, mesh, equations, initial_condition, boundary_conditions, source_terms, solver, boundaries, _node_coordinates, cell_ids, node_coordinates, inverse_jacobian, _neighbor_ids, neighbor_ids, orientation, surface_flux_values, u)

The best thing to do for a user would be to separate out the things that you need to track through, make them arguments to the function, and then simply Duplicate on those.

jacobian_enzyme_forward(semi::SemidiscretizationHyperbolic; N = pick_batchsize(semi))

Uses the right-hand side operator of the semidiscretization semi and forward mode automatic differentiation to compute the Jacobian J of the semidiscretization semi at state u0_ode.

Notes

Both jacobian_enzyme_forward(semi) and jacobian_enzyme_reverse(semi) depend on enzyme_rhs!.

jacobian_enzyme_forward(f!::F, x::AbstractArray; N = pick_batchsize(x)) where F <: Function

Uses the function f! and forward mode automatic differentiation to compute the Jacobian J

Examples

julia> x = -1:0.5:1;
julia> batch_size = 2;
julia> jacobian_enzyme_forward(TrixiEnzyme.upwind!, x, N=batch_size)
5×5 Matrix{Float64}:
 -0.2  -0.0  -0.0  -0.0   0.2
  0.2  -0.2  -0.0  -0.0  -0.0
 -0.0   0.2  -0.2  -0.0  -0.0
 -0.0  -0.0   0.2  -0.2  -0.0
 -0.0  -0.0  -0.0   0.2  -0.2

jacobian_enzyme_forward_closure(semi::SemidiscretizationHyperbolic; N = pick_batchsize(semi))

Same as jacobian_enzyme_forward but with closure

Notes

I resolved issues related to type instability caused by closures, which is a known limitation of Enzyme.

I utilized closures here because they simplify the reuse of memory buffers and temporary variables without the need for explicit storage. let blocks create a new hard scope and optionally introduce new local bindings.

jacobian_enzyme_reverse(semi::SemidiscretizationHyperbolic)

The jacobianenzymereverse reuses containers for shadow variables. This might provide beginners with a clear view of how Enzyme works in detail.

jacobian_enzyme_reverse_closure(semi::SemidiscretizationHyperbolic)

Same as jacobian_enzyme_reverse but with closure

pick_batchsize(x)
pick_batchsize(semi)

Return a reasonable batch size for batched differentiation.

Arguments

• x: AbstractArray
• semi: SemidiscretizationHyperbolic in Trixi.jl

Notes

Inspired by https://github.com/EnzymeAD/Enzyme.jl/pull/1545/files

Examples

julia> pick_batchsize(rand(3))
3

julia> pick_batchsize(rand(20))
11

upwind!(du, u, cache)

Vanilla upwind scheme