diff --git a/.github/workflows/Benchmarks.yml b/.github/workflows/Benchmarks.yml index 6b2247fdc..9fe1ba3a6 100644 --- a/.github/workflows/Benchmarks.yml +++ b/.github/workflows/Benchmarks.yml @@ -37,7 +37,7 @@ jobs: - uses: actions/checkout@v5 - uses: julia-actions/setup-julia@v2 with: - version: '1' + version: '1.11' arch: x64 - uses: julia-actions/cache@v2 - name: Run benchmark diff --git a/benchmarks/Project.toml b/benchmarks/Project.toml index 98c29a643..1888d7890 100644 --- a/benchmarks/Project.toml +++ b/benchmarks/Project.toml @@ -1,4 +1,8 @@ [deps] BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" QuantumToolbox = "6c2fb7c5-b903-41d2-bc5e-5a7c320b9fab" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" +SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" diff --git a/benchmarks/autodiff.jl b/benchmarks/autodiff.jl new file mode 100644 index 000000000..8303ed1de --- /dev/null +++ b/benchmarks/autodiff.jl @@ -0,0 +1,92 @@ +function benchmark_autodiff!(SUITE) + # Use harmonic oscillator system for both sesolve and mesolve + N = 20 + a = destroy(N) + ψ0 = fock(N, 0) + tlist = range(0, 40, 100) + + # ---- SESOLVE ---- + # For direct Forward differentiation + function my_f_sesolve_direct(p) + H = p[1] * a' * a + p[2] * (a + a') + sol = sesolve(H, ψ0, tlist, progress_bar = Val(false)) + return real(expect(a' * a, sol.states[end])) + end + + # For SciMLSensitivity.jl (reverse mode with Zygote and Enzyme) + coef_Δ(p, t) = p[1] + coef_F(p, t) = p[2] + H_evo = QobjEvo(a' * a, coef_Δ) + QobjEvo(a + a', coef_F) + + function my_f_sesolve(p) + sol = sesolve( + H_evo, + ψ0, + tlist, + progress_bar = Val(false), + params = p, + sensealg = BacksolveAdjoint(autojacvec = EnzymeVJP()), + ) + return real(expect(a' * a, sol.states[end])) + end + + # ---- MESOLVE ---- + # For direct Forward differentiation + function my_f_mesolve_direct(p) + H = p[1] * a' * a + p[2] * (a + a') + c_ops = [sqrt(p[3]) * a] + sol = mesolve(H, ψ0, tlist, c_ops, progress_bar = Val(false)) + return real(expect(a' * a, sol.states[end])) + end + + # For SciMLSensitivity.jl (reverse mode with Zygote and Enzyme) + coef_γ(p, t) = sqrt(p[3]) + c_ops = [QobjEvo(a, coef_γ)] + L = liouvillian(H_evo, c_ops) + + function my_f_mesolve(p) + sol = mesolve( + L, + ψ0, + tlist, + progress_bar = Val(false), + params = p, + sensealg = BacksolveAdjoint(autojacvec = EnzymeVJP()), + ) + return real(expect(a' * a, sol.states[end])) + end + + # Parameters for benchmarks + params_sesolve = [1.0, 1.0] + params_mesolve = [1.0, 1.0, 1.0] + + # Benchmark sesolve - Forward + SUITE["Autodiff"]["sesolve"]["Forward"] = @benchmarkable ForwardDiff.gradient($my_f_sesolve_direct, $params_sesolve) + + # Benchmark sesolve - Reverse (Zygote) + SUITE["Autodiff"]["sesolve"]["Reverse (Zygote)"] = @benchmarkable Zygote.gradient($my_f_sesolve, $params_sesolve) + + # Benchmark sesolve - Reverse (Enzyme) + SUITE["Autodiff"]["sesolve"]["Reverse (Enzyme)"] = @benchmarkable Enzyme.autodiff( + Enzyme.set_runtime_activity(Enzyme.Reverse), + Const($my_f_sesolve), + Active, + Duplicated($params_sesolve, dparams_sesolve), + ) setup=(dparams_sesolve = Enzyme.make_zero($params_sesolve)) + + # Benchmark mesolve - Forward + SUITE["Autodiff"]["mesolve"]["Forward"] = @benchmarkable ForwardDiff.gradient($my_f_mesolve_direct, $params_mesolve) + + # Benchmark mesolve - Reverse (Zygote) + SUITE["Autodiff"]["mesolve"]["Reverse (Zygote)"] = @benchmarkable Zygote.gradient($my_f_mesolve, $params_mesolve) + + # Benchmark mesolve - Reverse (Enzyme) + SUITE["Autodiff"]["mesolve"]["Reverse (Enzyme)"] = @benchmarkable Enzyme.autodiff( + Enzyme.set_runtime_activity(Enzyme.Reverse), + Const($my_f_mesolve), + Active, + Duplicated($params_mesolve, dparams_mesolve), + ) setup=(dparams_mesolve = Enzyme.make_zero($params_mesolve)) + + return nothing +end diff --git a/benchmarks/runbenchmarks.jl b/benchmarks/runbenchmarks.jl index e0252073d..ea3cc9ddb 100644 --- a/benchmarks/runbenchmarks.jl +++ b/benchmarks/runbenchmarks.jl @@ -3,6 +3,10 @@ using LinearAlgebra using SparseArrays using QuantumToolbox using SciMLBase: EnsembleSerial, EnsembleThreads +using ForwardDiff +using Zygote +using Enzyme: Enzyme, Const, Active, Duplicated +using SciMLSensitivity: BacksolveAdjoint, EnzymeVJP BLAS.set_num_threads(1) @@ -14,6 +18,7 @@ include("dynamical_shifted_fock.jl") include("eigenvalues.jl") include("steadystate.jl") include("timeevolution.jl") +include("autodiff.jl") benchmark_correlations_and_spectrum!(SUITE) benchmark_dfd!(SUITE) @@ -21,6 +26,7 @@ benchmark_dsf!(SUITE) benchmark_eigenvalues!(SUITE) benchmark_steadystate!(SUITE) benchmark_timeevolution!(SUITE) +benchmark_autodiff!(SUITE) BenchmarkTools.tune!(SUITE) results = BenchmarkTools.run(SUITE, verbose = true)