add enzyme and mooncake to benchmark

CDCgov · Dec 13, 2024 · 4f3d624 · 4f3d624
1 parent 12a21c6
commit 4f3d624
Show file tree

Hide file tree

Showing 7 changed files with 283 additions and 183 deletions.
diff --git a/benchmark/Manifest.toml b/benchmark/Manifest.toml
diff --git a/benchmark/Project.toml b/benchmark/Project.toml
@@ -1,9 +1,12 @@
 [deps]
+ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
 BenchmarkCI = "20533458-34a3-403d-a444-e18f38190b5b"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 EpiAware = "b2eeebe4-5992-4301-9193-7ebc9f62c855"
+Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6"
 PkgBenchmark = "32113eaa-f34f-5b0d-bd6c-c81e245fc73d"
 Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
 TuringBenchmarking = "0db1332d-5c25-4deb-809f-459bc696f94f"
diff --git a/benchmark/bench/EpiAwareUtils/EpiAwareUtils.jl b/benchmark/bench/EpiAwareUtils/EpiAwareUtils.jl
@@ -1,6 +1,9 @@
 module BenchEpiAwareUtils
 
 using BenchmarkTools, EpiAware.EpiAwareUtils
+using ADTypes, Mooncake, Enzyme
+
+Enzyme.API.runtimeActivity!(true)
 
 suite = BenchmarkGroup()
 

diff --git a/benchmark/bench/EpiInfModels/EpiInfModels.jl b/benchmark/bench/EpiInfModels/EpiInfModels.jl
@@ -1,6 +1,10 @@
 module BenchEpiInfModels
 
 using BenchmarkTools, TuringBenchmarking, EpiAware, Distributions
+using ADTypes, Mooncake, Enzyme
+
+Enzyme.API.runtimeActivity!(true)
+
 suite = BenchmarkGroup()
 
 include("../../make_epiaware_suite.jl")

diff --git a/benchmark/bench/EpiLatentModels/EpiLatentModels.jl b/benchmark/bench/EpiLatentModels/EpiLatentModels.jl
@@ -1,6 +1,10 @@
 module BenchEpiLatentModels
 
 using BenchmarkTools, TuringBenchmarking, EpiAware, DynamicPPL
+using ADTypes, Mooncake, Enzyme
+
+Enzyme.API.runtimeActivity!(true)
+
 suite = BenchmarkGroup()
 
 include("../../make_epiaware_suite.jl")

diff --git a/benchmark/bench/EpiObsModels/EpiObsModels.jl b/benchmark/bench/EpiObsModels/EpiObsModels.jl
@@ -1,6 +1,10 @@
 module BenchEpiObsModels
 
 using BenchmarkTools, TuringBenchmarking, EpiAware, DynamicPPL
+using ADTypes, Mooncake, Enzyme
+
+Enzyme.API.runtimeActivity!(true)
+
 suite = BenchmarkGroup()
 
 include("../../make_epiaware_suite.jl")

diff --git a/benchmark/make_epiaware_suite.jl b/benchmark/make_epiaware_suite.jl
@@ -2,7 +2,11 @@
    A custom wrapper for the `TuringBenchmarking.make_turing_suite` that adds EpiAware specific defaults.
 """
 function make_epiaware_suite(model; check = true,
-        adbackends = [:forwarddiff, :reversediff, :reversediff_compiled])
+        adbackends = [
+            :forwarddiff, :reversediff, :reversediff_compiled,
+            ADTypes.AutoMooncake(config = nothing),
+            ADTypes.AutoEnzyme()
+        ])
     suite = prefix_warnings(
         () -> TuringBenchmarking.make_turing_suite(
             model; check = check, adbackends = adbackends),