Add Enzyme benchmarks

Author: albertomercurio

benchmarks/Project.toml

@@ -1,5 +1,6 @@
 [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"

benchmarks/autodiff.jl

@@ -60,17 +60,37 @@ function benchmark_autodiff!(SUITE)
     params_sesolve = [1.0, 1.0]
     params_mesolve = [1.0, 1.0, 1.0]
 
+    # Pre-allocate gradient arrays for Enzyme
+    dparams_sesolve = Enzyme.make_zero(params_sesolve)
+    dparams_mesolve = Enzyme.make_zero(params_mesolve)
+
     # Benchmark sesolve - Forward
     SUITE["Autodiff"]["sesolve"]["Forward"] = @benchmarkable ForwardDiff.gradient($my_f_sesolve_direct, $params_sesolve)
 
     # Benchmark sesolve - Reverse (Zygote)
-    SUITE["Autodiff"]["sesolve"]["Reverse"] = @benchmarkable Zygote.gradient($my_f_sesolve, $params_sesolve)
+    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),
+    )
 
     # Benchmark mesolve - Forward
     SUITE["Autodiff"]["mesolve"]["Forward"] = @benchmarkable ForwardDiff.gradient($my_f_mesolve_direct, $params_mesolve)
 
     # Benchmark mesolve - Reverse (Zygote)
-    SUITE["Autodiff"]["mesolve"]["Reverse"] = @benchmarkable Zygote.gradient($my_f_mesolve, $params_mesolve)
+    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),
+    )
 
     return nothing
 end

benchmarks/runbenchmarks.jl

@@ -5,6 +5,7 @@ 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)