Performance of high order derivatives in Enzyme is slower than finite difference
- 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
Assessment
This issue has not been assessed yet.