Skip to content

Commit

Permalink
chore:
Browse files Browse the repository at this point in the history
  • Loading branch information
junyixu committed Aug 24, 2024
1 parent aa30582 commit 086fe0c
Show file tree
Hide file tree
Showing 4 changed files with 25 additions and 1 deletion.
File renamed without changes
13 changes: 12 additions & 1 deletion docs/src/index.md
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,17 @@ TrixiEnzyme is not a registered Julia package, and it can be installed by runnin
] add https://github.com/junyixu/TrixiEnzyme.jl.git
```

## Notes about Enzyme

There's some issues with `Enzyme.make_zero!` right now for this use case, see https://github.com/EnzymeAD/Enzyme.jl/issues/1661
One need to be careful with a vanilla closure outside Enzyme.
If one writes to caches and expect to differentiate through them, then the closure should be duplicated for handling the derivative of those values.
If you want to track derivatives through arrays that are enclosed, you have to duplicate the array to have shadow memory for its differentiation.
if you want to track derivatives through arrays that are enclosed, you have to duplicate the array to have shadow memory for its differentiation
So if you only have the original memory, you cannot do the differentiation since you don't have a place to store the extra values. In a simplified sense, a `Dual{Float64}` is 128 bits, `Float64` is 64 bits, so if you're writing to a buffer of 5 `Float64` numbers, you need 5*2*64 bits of space to keep a dual number, which you don't have
So 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.
This is how [TrixiEnzyme.jacobian_enzyme_forward](https://junyixu.github.io/TrixiEnzyme.jl/dev/api.html#TrixiEnzyme.jacobian_enzyme_forward) works.

## Configuring Batch Size

To utilize `Enzyme.BatchDuplicated`, one can create a tuple containing duals (or shadows).
Expand Down Expand Up @@ -45,5 +56,5 @@ julia> @time jacobian_enzyme_forward(TrixiEnzyme.upwind!, x);
```

Benchmark for a 401x401 Jacobian of `TrixiEnzyme.upwind!` (Lower is better):
![upwind benchmark](../fig/upwind_benchmark.png)
![upwind benchmark](./img/upwind_benchmark.png)
`Enyme(@batch)` means applying `Polyester.@batch` to `middlebatches`.
12 changes: 12 additions & 0 deletions src/jacobian.jl
Original file line number Diff line number Diff line change
@@ -1,5 +1,10 @@
Enzyme.API.runtimeActivity!(true)

```
jacobian_enzyme_forward_closure(semi::SemidiscretizationHyperbolic)
Same as jacobian_enzyme_forward but with closure.
```
function jacobian_enzyme_forward_closure(semi)
t0 = zero(real(semi))
u_ode = compute_coefficients(t0, semi)
Expand Down Expand Up @@ -32,7 +37,14 @@ function jacobian_enzyme_forward_closure(semi)
return dys
end

```
jacobian_enzyme_reverse_closure(semi::SemidiscretizationHyperbolic)
Same as jacobian_enzyme_reverse but with closure
!!! warning
Enzyme.jl does not play well with Polyester.jl and there are no plans to fix this soon.
```
function jacobian_enzyme_reverse_closure(semi)
t0 = zero(real(semi))
u_ode = compute_coefficients(t0, semi)
Expand Down
1 change: 1 addition & 0 deletions test/SemiTest.jl
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
module SemiTest
using Test
using TrixiEnzyme
using TrixiEnzyme: LinearScalarAdvectionEquation1D, DGSEM, TreeMesh, SVector

# %%
# equation with a advection_velocity of `1`.
Expand Down

0 comments on commit 086fe0c

Please sign in to comment.