EnzymeAD / EnzymeAD/Enzyme.jl

Performance of high order derivatives in Enzyme is slower than finite difference

Open
#2,252 3 comments 0 reactions 0 assignees View on GitHub
Dominant language
Julia
Stars
586
Forks
108
Avg merge
1d 5h
Merged PRs (30d)
44

Description

Thank you for adding the capability of computing high order derivatives with Enzyme in https://github.com/EnzymeAD/Enzyme.jl/pull/2161!

I benchmarked the performance of Enzyme against finite difference method. For orders greater than 2, I see that finite difference is faster. It is 1.7 times faster for order 3 and 2.4 times faster for order 4. I am sharing my benchmarking code for orders 3 and 4 in case it leads to any ideas for improvement. The mean performance of the third order finite difference method is 8 ns, while that of Enzyme is 14 ns. For the fourth order, finite difference is at 11 ns and Enzyme is at 27 ns.

Benchmarking code for third order

```julia
using BenchmarkTools: @benchmark
using StaticArrays
using Enzyme

# Flux function to be differentiation

@inline function flux(u)
rho, rho_v1, rho_v2, rho_e = u
gamma = 1.4
v1 = rho_v1 / rho
v2 = rho_v2 / rho
p = (gamma - 1) * (rho_e - 0.5f0 * (rho_v1 * v1 + rho_v2 * v2))
f1 = rho_v1
f2 = rho_v1 * v1 + p
f3 = rho_v1 * v2
f4 = (rho_e + p) * v1
return SVector(f1, f2, f3, f4)
end

# Third order finite difference derivative

function third_derivative_fd(u, du, ddu, dddu)
factor = 0.5
df = factor * (flux(u + 2.0 * du + 2.0 * ddu + 4.0/3.0 * dddu)
- 2.0 * flux(u + du + 0.5 * ddu + (1.0/6.0) * dddu)
+ 2.0 * flux(u - du + 0.5 * ddu - (1.0/6.0) * dddu)
- flux(u - 2.0 * du + 2.0 * ddu - 4.0/3.0 * dddu))
return df
end

# AD to compute derivatives

dg_ad(x, dx) = autodiff(Forward, flux, DuplicatedNoNeed(x, dx))[1]
ddg_ad(x, dx, ddx) = autodiff(Forward, dg_ad, DuplicatedNoNeed(x, dx),
DuplicatedNoNeed(dx, ddx))[1]
dddg_ad(x, dx, ddx, dddx) = autodiff(Forward, ddg_ad, DuplicatedNoNeed(x, dx),
DuplicatedNoNeed(dx, ddx), DuplicatedNoNeed(ddx, dddx))[1]

# Random inputs

u = SVector(1.0, -0.1, 0.2, 2.0)
du, ddu, dddu, ddddu = (1e-3*SVector(rand(4)...) for _ in 1:4)

@info "Third derivative"
@info "FD"
display(@benchmark third_derivative_fd($u, $du, $ddu, $dddu))

@info "Enzyme"

display(@benchmark dddg_ad($u, $du, $ddu, $dddu))
```

Benchmarking code for fourth order

```julia
using BenchmarkTools: @benchmark
using StaticArrays
using Enzyme

# Flux function to be differentiation

@inline function flux(u)
rho, rho_v1, rho_v2, rho_e = u
gamma = 1.4
v1 = rho_v1 / rho
v2 = rho_v2 / rho
p = (gamma - 1) * (rho_e - 0.5f0 * (rho_v1 * v1 + rho_v2 * v2))
f1 = rho_v1
f2 = rho_v1 * v1 + p
f3 = rho_v1 * v2
f4 = (rho_e + p) * v1
return SVector(f1, f2, f3, f4)
end

# Fourth order finite difference derivative

function fourth_derivative_fd(u, du, ddu, dddu, ddddu)
df = (
flux(u + 2.0 * du + 2.0 * ddu + 4.0/3.0 * dddu + 2.0/3.0 * ddddu)
- 4.0 * flux(u + du + 0.5 * ddu + 1.0/6.0 * dddu + 1.0/24.0 * ddddu)
+ 6.0 * flux(u)
- 4.0 * flux(u - du + 0.5 * ddu - 1.0/6.0 * dddu + 1.0/24.0 * ddddu)
+ flux(u - 2.0 * du + 2.0 * ddu - 4.0/3.0 * dddu + 2.0/3.0 * ddddu)
)
return df
end

# AD to compute derivatives

dg_ad(x, dx) = autodiff(Forward, flux, DuplicatedNoNeed(x, dx))[1]
ddg_ad(x, dx, ddx) = autodiff(Forward, dg_ad, DuplicatedNoNeed(x, dx),
DuplicatedNoNeed(dx, ddx))[1]
dddg_ad(x, dx, ddx, dddx) = autodiff(Forward, ddg_ad, DuplicatedNoNeed(x, dx),
DuplicatedNoNeed(dx, ddx), DuplicatedNoNeed(ddx, dddx))[1]
ddddg_ad(x, dx, ddx, dddx, ddddx) = autodiff(Forward, dddg_ad, DuplicatedNoNeed(x, dx),
DuplicatedNoNeed(dx, ddx),
DuplicatedNoNeed(ddx, dddx),
DuplicatedNoNeed(dddx, ddddx))[1]

# Random inputs

u = SVector(1.0, -0.1, 0.2, 2.0)
du, ddu, dddu, ddddu = (1e-3*SVector(rand(4)...) for _ in 1:4)

@info "Fourth derivative"
@info "FD"
display(@benchmark fourth_derivative_fd($u, $du, $ddu, $dddu, $dddu))

@info "Enzyme"

display(@benchmark ddddg_ad($u, $du, $ddu, $dddu, $dddu))
```

Benchmarking results for third order

```julia-repl
[ Info: Third derivative
[ Info: FD
BenchmarkTools.Trial: 10000 samples with 999 evaluations.
Range (min … max): 7.924 ns … 20.604 ns ┊ GC (min … max): 0.00% … 0.00%
Time (median): 8.008 ns ┊ GC (median): 0.00%
Time (mean ± σ): 8.140 ns ± 0.710 ns ┊ GC (mean ± σ): 0.00% ± 0.00%

██▅▃▃▂ ▃ ▁▁ ▂
██████▇▆▆▅▆█▆▆▆▅▃█▇▄▅▅▄▅▄▅▄▃▅▅▁▅▄▄▆▄▃▅▄▃▃▁▄▅▅▃▄▃▄▄▄▅▆▅▆▇██ █
7.92 ns Histogram: log(frequency) by time 11.2 ns <

Memory estimate: 0 bytes, allocs estimate: 0.
[ Info: Enzyme
BenchmarkTools.Trial: 10000 samples with 999 evaluations.
Range (min … max): 13.639 ns … 26.652 ns ┊ GC (min … max): 0.00% … 0.00%
Time (median): 13.722 ns ┊ GC (median): 0.00%
Time (mean ± σ): 14.037 ns ± 0.980 ns ┊ GC (mean ± σ): 0.00% ± 0.00%

█▅▃ ▂ ▁▅▂▁ ▁
███▇█▇▆▆▃▆████▇▅▄▅▇▅▅▆▆▅▄▅▄▁▄▄▅▅▅▄▆▄▆▃▅▃▄▅▄▅▅▆▅▆▄▅▅▄▅▄▅▅▄▆█ █
13.6 ns Histogram: log(frequency) by time 19.1 ns <

```

Benchmarking results for fourth order

```julia-repl
[ Info: Fourth derivative
[ Info: FD
BenchmarkTools.Trial: 10000 samples with 999 evaluations.
Range (min … max): 10.927 ns … 83.625 ns ┊ GC (min … max): 0.00% … 0.00%
Time (median): 12.304 ns ┊ GC (median): 0.00%
Time (mean ± σ): 12.644 ns ± 2.999 ns ┊ GC (mean ± σ): 0.00% ± 0.00%

▅▃█▂ █▆▃▃▂▁▁▁ ▁▁▁▂▁▁▁▂▁▂▄▂▁▁ ▂
███████████████████████████████▅▅▄▄▅▄▅▅▄▅▄▄▄▄▄▄▅▃▄▁▄▁▄▆▅▅▆▅ █
10.9 ns Histogram: log(frequency) by time 20.5 ns <

Memory estimate: 0 bytes, allocs estimate: 0.
[ Info: Enzyme
BenchmarkTools.Trial: 10000 samples with 996 evaluations.
Range (min … max): 26.230 ns … 43.257 ns ┊ GC (min … max): 0.00% … 0.00%
Time (median): 26.397 ns ┊ GC (median): 0.00%
Time (mean ± σ): 26.774 ns ± 1.646 ns ┊ GC (mean ± σ): 0.00% ± 0.00%

█▆ ▁ ▁
███▆█▇▆▅▅▆▆▆▆▆▅▅▄▅▆▇▇▆▅▆▅▅▅▅▆▅▅▄▄▄▄▅▅▅▄▅▅▄▆▅▅▆▅▄▄▅▅▅▄▅▄▅▃▄▅ █
26.2 ns Histogram: log(frequency) by time 36.4 ns <

Memory estimate: 0 bytes, allocs estimate: 0.

```

These results have been generated with Enzyme `v0.13.24` using Apple M3 Pro on Julia 1.11.2.

versioninfo()

```julia-repl
Julia Version 1.11.2
Commit 5e9a32e7af2 (2024-12-01 20:02 UTC)
Build Info:
Official https://julialang.org/ release
Platform Info:
OS: macOS (arm64-apple-darwin24.0.0)
CPU: 12 × Apple M3 Pro
WORD_SIZE: 64
LLVM: libLLVM-16.0.6 (ORCJIT, apple-m2)
Threads: 1 default, 0 interactive, 1 GC (on 6 virtual cores)
```

[Here](https://gist.github.com/Arpit-Babbar/f7173c8ed83cbd2a910dbe4fb5bf3210) is a gist that checks that the above third and fourth order derivative computations are correct by using a polynomial test case. [This gist](https://gist.github.com/Arpit-Babbar/53591d5d4047646553ec076f1b645247) contains computation and benchmarking of all derivatives up to four.

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.