Skip to content

Support Enzyme and Mooncake in AD backends #291

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 21 commits into
base: master
Choose a base branch
from

Conversation

ErikQQY
Copy link
Member

@ErikQQY ErikQQY commented Mar 19, 2025

Part of #288

@ErikQQY
Copy link
Member Author

ErikQQY commented Mar 19, 2025

Seems there are still some thing to do:

julia> @benchmark sol = solve(prob, MIRK4(; jac_alg = jac_alg_forwarddiff), dt = 0.05)
BenchmarkTools.Trial: 7002 samples with 1 evaluation per sample.
 Range (min  max):  378.708 μs  444.849 ms  ┊ GC (min  max):  0.00%  97.30%
 Time  (median):     433.500 μs               ┊ GC (median):     0.00%
 Time  (mean ± σ):   767.981 μs ±   5.955 ms  ┊ GC (mean ± σ):  21.83% ± 10.28%

  █▃                                                            ▁
  ███▇▆▆▅▅▅▅▅▅▄▄▄▁▄▁▃▅▁▁▄▁▄▁▃▃▃▃▁▄▄▁▃▁▁▃▃▄▁▁▁▁▄▃▃▄▄▅▆▅▅▄▄▁▃▁▃▅▅ █
  379 μs        Histogram: log(frequency) by time       7.72 ms <

 Memory estimate: 1.46 MiB, allocs estimate: 27500.

julia> @benchmark sol = solve(prob, MIRK4(; jac_alg = jac_alg_enzyme), dt = 0.05)
BenchmarkTools.Trial: 3252 samples with 1 evaluation per sample.
 Range (min  max):  1.227 ms  32.662 ms  ┊ GC (min  max):  0.00%  95.57%
 Time  (median):     1.344 ms              ┊ GC (median):     0.00%
 Time  (mean ± σ):   1.536 ms ±  1.057 ms  ┊ GC (mean ± σ):  11.47% ± 14.84%

  ▄█▅▃▁                                                       
  █████▇▅▅▄▅▃▃▄▃▃▃▁▃▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▅▆▆▅▇▇▇▇▆▆▅▆▆▅▃▅ █
  1.23 ms      Histogram: log(frequency) by time      5.3 ms <

 Memory estimate: 2.43 MiB, allocs estimate: 27517.

Copy link
Contributor

github-actions bot commented Mar 19, 2025

Benchmark Results

master c7c3c64... master / c7c3c64...
Simple Pendulum/IIP/BoundaryValueDiffEqMIRK.MIRK2() 7.79 ± 0.19 ms 7.69 ± 0.16 ms 1.01
Simple Pendulum/IIP/BoundaryValueDiffEqMIRK.MIRK3() 2.87 ± 0.15 ms 2.85 ± 0.13 ms 1.01
Simple Pendulum/IIP/BoundaryValueDiffEqMIRK.MIRK4() 1.01 ± 0.082 ms 0.994 ± 0.06 ms 1.01
Simple Pendulum/IIP/BoundaryValueDiffEqMIRK.MIRK5() 2.61 ± 0.44 ms 2.57 ± 0.25 ms 1.02
Simple Pendulum/IIP/BoundaryValueDiffEqMIRK.MIRK6() 1.19 ± 0.14 ms 1.17 ± 0.1 ms 1.02
Simple Pendulum/IIP/MultipleShooting(10, Tsit5; grid_coarsening = false) 1.85 ± 0.66 ms 1.83 ± 0.69 ms 1.01
Simple Pendulum/IIP/MultipleShooting(10, Tsit5; grid_coarsening = true) 3.14 ± 1 ms 3.09 ± 1 ms 1.02
Simple Pendulum/IIP/MultipleShooting(100, Tsit5; grid_coarsening = false) 0.0689 ± 0.021 s 0.0698 ± 0.021 s 0.987
Simple Pendulum/IIP/MultipleShooting(100, Tsit5; grid_coarsening = true) 0.083 ± 0.021 s 0.0825 ± 0.024 s 1.01
Simple Pendulum/IIP/Shooting(Tsit5()) 0.252 ± 0.078 ms 0.247 ± 0.078 ms 1.02
Simple Pendulum/OOP/BoundaryValueDiffEqMIRK.MIRK2() 8.83 ± 0.43 ms 9.02 ± 0.48 ms 0.979
Simple Pendulum/OOP/BoundaryValueDiffEqMIRK.MIRK3() 3.3 ± 0.15 ms 3.33 ± 0.22 ms 0.991
Simple Pendulum/OOP/BoundaryValueDiffEqMIRK.MIRK4() 1.17 ± 0.089 ms 1.16 ± 0.088 ms 1.01
Simple Pendulum/OOP/BoundaryValueDiffEqMIRK.MIRK5() 2.94 ± 0.31 ms 2.97 ± 0.31 ms 0.99
Simple Pendulum/OOP/BoundaryValueDiffEqMIRK.MIRK6() 1.39 ± 0.12 ms 1.39 ± 0.11 ms 1
Simple Pendulum/OOP/MultipleShooting(10, Tsit5; grid_coarsening = false) 3.58 ± 3 ms 3.6 ± 3 ms 0.995
Simple Pendulum/OOP/MultipleShooting(10, Tsit5; grid_coarsening = true) 6.03 ± 5 ms 6.04 ± 5.2 ms 0.999
Simple Pendulum/OOP/MultipleShooting(100, Tsit5; grid_coarsening = false) 0.124 ± 0.022 s 0.137 ± 0.026 s 0.906
Simple Pendulum/OOP/MultipleShooting(100, Tsit5; grid_coarsening = true) 0.156 ± 0.028 s 0.156 ± 0.026 s 1
Simple Pendulum/OOP/Shooting(Tsit5()) 0.687 ± 0.052 ms 0.692 ± 0.052 ms 0.993
time_to_load 4.96 ± 0.073 s 5.09 ± 0.056 s 0.975

Benchmark Plots

A plot of the benchmark results have been uploaded as an artifact to the workflow run for this PR.
Go to "Actions"->"Benchmark a pull request"->[the most recent run]->"Artifacts" (at the bottom).

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
None yet
Projects
None yet
Development

Successfully merging this pull request may close these issues.

1 participant